drop schema on teardown

pgdata/__init__.py

@@ -53,6 +53,8 @@ def drop_db(url):
            FROM pg_database
            WHERE datname = '{db_name}'""".format(db_name=db_name)
     if db.query(q).fetchone():
+        # DROP DATABASE must be run outside of a transaction
         conn = db.engine.connect()
+        conn.execute("commit")
         conn.execute("DROP DATABASE "+db_name)
         conn.close()

tests/test_pg2ogr.py

@@ -5,13 +5,13 @@
 
 import fiona
 
-from pgdata import connect
+import pgdata
 
 
 URL = 'postgresql://postgres:postgres@localhost:5432/pgdata'
-DB = connect(URL)
+DB = pgdata.connect(URL)
 DB.execute('CREATE SCHEMA IF NOT EXISTS pgdata')
-DB1 = connect(URL)
+
 
 DATA = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data')
 AIRPORTS = os.path.join(DATA, 'bc_airports.json')
@@ -21,34 +21,31 @@ class ogrpg(unittest.TestCase):
     def setUp(self):
         self.tempdir = tempfile.mkdtemp()
 
-    def tearDown(self):
-        shutil.rmtree(self.tempdir)
-
     def test_ogr2pg(self):
-        db = DB1
+        db = DB
         db.ogr2pg(AIRPORTS, in_layer='bc_airports', out_layer='bc_airports',
                   schema='pgdata')
         airports = db['pgdata.bc_airports']
         assert 'physical_address' in airports.columns
         assert sum(1 for _ in airports.all()) == 425
 
     def test_pg2geojson(self):
-        db = DB1
+        db = DB
         db.pg2ogr(sql='SELECT * FROM pgdata.bc_airports LIMIT 10', driver='GeoJSON',
                   outfile=os.path.join(self.tempdir, 'test_dump.json'))
         c = fiona.open(os.path.join(self.tempdir, 'test_dump.json'), 'r')
         assert len(c) == 10
 
     def test_pg2gpkg(self):
-        db = DB1
+        db = DB
         db.pg2ogr(sql='SELECT * FROM pgdata.bc_airports LIMIT 10', driver='GPKG',
                   outfile=os.path.join(self.tempdir, 'test_dump.gpkg'),
                   outlayer='bc_airports')
         c = fiona.open(os.path.join(self.tempdir, 'test_dump.gpkg'), 'r')
         assert len(c) == 10
 
     def test_pg2gpkg_update(self):
-        db = DB1
+        db = DB
         db.pg2ogr(sql='SELECT * FROM pgdata.bc_airports LIMIT 10', driver='GPKG',
                   outfile=os.path.join(self.tempdir, 'test_dump.gpkg'),
                   outlayer='bc_airports')
@@ -59,7 +56,7 @@ def test_pg2gpkg_update(self):
         assert len(layers) == 2
 
     def test_pg2ogr_append(self):
-        db = DB1
+        db = DB
         db.pg2ogr(sql='SELECT * FROM pgdata.bc_airports LIMIT 10', driver='GPKG',
                   outfile=os.path.join(self.tempdir, 'test_dump.gpkg'),
                   outlayer='bc_airports')
@@ -68,3 +65,10 @@ def test_pg2ogr_append(self):
                   outlayer='bc_airports', append=True)
         c = fiona.open(os.path.join(self.tempdir, 'test_dump.gpkg'), 'r')
         assert len(c) == 20
+
+    def tearDown(self):
+        shutil.rmtree(self.tempdir)
+
+
+def test_tearDown():
+    DB.drop_schema('pgdata', cascade=True)