diff --git a/setup.py b/setup.py
index ad8c7da..73b4a63 100644
--- a/setup.py
+++ b/setup.py
@@ -36,6 +36,7 @@ setup(
         "PyYAML >= 3.10"
     ] + py26_dependency,
     tests_require=[],
+    test_suite='test',
     entry_points={
         'console_scripts': [
             'datafreeze = dataset.freeze.app:main',
diff --git a/test/test_persistence.py b/test/test_persistence.py
index 1ae95fc..998b2e4 100644
--- a/test/test_persistence.py
+++ b/test/test_persistence.py
@@ -2,10 +2,12 @@ import os
 import unittest
 from datetime import datetime
 
+from sqlalchemy.exc import IntegrityError
+
 from dataset import connect
 from dataset.util import DatasetException
-from sample_data import TEST_DATA, TEST_CITY_1
-from sqlalchemy.exc import IntegrityError
+
+from .sample_data import TEST_DATA, TEST_CITY_1
 
 
 class DatabaseTestCase(unittest.TestCase):
@@ -241,6 +243,3 @@ class TableTestCase(unittest.TestCase):
         assert 'foo' in tbl.table.c, tbl.table.c
         assert FLOAT == type(tbl.table.c['foo'].type), tbl.table.c['foo'].type
         assert 'foo' in tbl.columns, tbl.columns
-
-if __name__ == '__main__':
-    unittest.main()