diff --git a/test/test_persistence.py b/test/test_persistence.py index 998b2e4..3feee9d 100644 --- a/test/test_persistence.py +++ b/test/test_persistence.py @@ -13,6 +13,7 @@ from .sample_data import TEST_DATA, TEST_CITY_1 class DatabaseTestCase(unittest.TestCase): def setUp(self): + self.old_db_url = os.environ.get('DATABASE_URL') os.environ['DATABASE_URL'] = 'sqlite:///:memory:' self.db = connect('sqlite:///:memory:') self.tbl = self.db['weather'] @@ -21,7 +22,10 @@ class DatabaseTestCase(unittest.TestCase): def tearDown(self): # ensure env variable was unset - del os.environ['DATABASE_URL'] + if self.old_db_url is None: + del os.environ['DATABASE_URL'] + else: + os.environ['DATABASE_URL'] = self.old_db_url def test_valid_database_url(self): assert self.db.url, os.environ['DATABASE_URL']