Replace parameterized tests with per-connection type tests with transactions.

Author: JacobHayes

tests/test_export.py

@@ -8,11 +8,6 @@
 
 Base = declarative_base()
 engine = sa.create_engine('postgresql:///copy-test')
-connection_types = {
-    'engine': lambda session: session.connection().engine,
-    'connection': lambda session: session.connection(),
-    'raw_connection': lambda session: session.connection().connection,
-}
 
 class Album(Base):
     __tablename__ = 'album'
@@ -44,79 +39,76 @@ def objects(session):
     finally:
         engine.execute(Album.__table__.delete())
 
-@pytest.mark.parametrize("conn_type", connection_types.values())
 class TestCopyTo:
 
-    def test_copy_query(self, session, objects, conn_type):
+    def test_copy_query(self, session, objects):
         sio = io.StringIO()
-        copy_to(session.query(Album), sio, conn_type(session))
+        copy_to(session.query(Album), sio, session.connection().engine)
         lines = sio.getvalue().strip().split('\n')
         assert len(lines) == 3
         assert lines[0].split('\t') == [str(objects[0].id), objects[0].name]
 
-    def test_copy_table(self, session, objects, conn_type):
+    def test_copy_table(self, session, objects):
         sio = io.StringIO()
-        copy_to(Album.__table__.select(), sio, conn_type(session))
+        copy_to(Album.__table__.select(), sio, session.connection().engine)
         lines = sio.getvalue().strip().split('\n')
         assert len(lines) == 3
         assert lines[0].split('\t') == [str(objects[0].id), objects[0].name]
 
-    def test_copy_csv(self, session, objects, conn_type):
+    def test_copy_csv(self, session, objects):
         sio = io.StringIO()
         flags = {'format': 'csv', 'header': True}
-        copy_to(session.query(Album), sio, conn_type(session), **flags)
+        copy_to(session.query(Album), sio, session.connection().engine, **flags)
         lines = sio.getvalue().strip().split('\n')
         assert len(lines) == 4
         assert lines[0].split(',') == ['aid', 'name']
         assert lines[1].split(',') == [str(objects[0].id), objects[0].name]
 
-@pytest.mark.parametrize("conn_type", connection_types.values())
 class TestCopyRename:
 
-    def test_rename_model(self, session, objects, conn_type):
+    def test_rename_model(self, session, objects):
         sio = io.StringIO()
         flags = {'format': 'csv', 'header': True}
         query = relabel_query(session.query(Album))
-        copy_to(query, sio, conn_type(session), **flags)
+        copy_to(query, sio, session.connection().engine, **flags)
         lines = sio.getvalue().strip().split('\n')
         assert len(lines) == 4
         assert lines[0].split(',') == ['id', 'name']
         assert lines[1].split(',') == [str(objects[0].id), objects[0].name]
 
-    def test_rename_columns(self, session, objects, conn_type):
+    def test_rename_columns(self, session, objects):
         sio = io.StringIO()
         flags = {'format': 'csv', 'header': True}
         query = relabel_query(session.query(Album.id, Album.name.label('title')))
-        copy_to(query, sio, conn_type(session), **flags)
+        copy_to(query, sio, session.connection().engine, **flags)
         lines = sio.getvalue().strip().split('\n')
         assert len(lines) == 4
         assert lines[0].split(',') == ['id', 'title']
         assert lines[1].split(',') == [str(objects[0].id), objects[0].name]
 
-@pytest.mark.parametrize("conn_type", connection_types.values())
 class TestCopyFrom:
 
-    def test_copy_model(self, session, objects, conn_type):
+    def test_copy_model(self, session, objects):
         sio = io.StringIO()
         sio.write(u'\t'.join(['4', 'The Works']))
         sio.seek(0)
-        copy_from(sio, Album, conn_type(session))
+        copy_from(sio, Album, session.connection().engine)
         assert session.query(Album).count() == len(objects) + 1
         row = session.query(Album).filter_by(id=4).first()
         assert row.id == 4
         assert row.name == 'The Works'
 
-    def test_copy_table(self, session, objects, conn_type):
+    def test_copy_table(self, session, objects):
         sio = io.StringIO()
         sio.write(u'\t'.join(['4', 'The Works']))
         sio.seek(0)
-        copy_from(sio, Album.__table__, conn_type(session))
+        copy_from(sio, Album.__table__, session.connection().engine)
         assert session.query(Album).count() == len(objects) + 1
         row = session.query(Album).filter_by(id=4).first()
         assert row.id == 4
         assert row.name == 'The Works'
 
-    def test_copy_csv(self, session, objects, conn_type):
+    def test_copy_csv(self, session, objects):
         sio = io.StringIO()
         sio.write(
             u'\n'.join([
@@ -126,18 +118,67 @@ def test_copy_csv(self, session, objects, conn_type):
         )
         sio.seek(0)
         flags = {'format': 'csv', 'header': True}
-        copy_from(sio, Album, conn_type(session), **flags)
+        copy_from(sio, Album, session.connection().engine, **flags)
         assert session.query(Album).count() == len(objects) + 1
         row = session.query(Album).filter_by(id=4).first()
         assert row.id == 4
         assert row.name == 'The Works'
 
-    def test_copy_columns(self, session, objects, conn_type):
+    def test_copy_columns(self, session, objects):
         sio = io.StringIO()
         sio.write(u'\t'.join(['The Works', '4']))
         sio.seek(0)
-        copy_from(sio, Album, conn_type(session), columns=('name', 'aid'))
+        copy_from(sio, Album, session.connection().engine, columns=('name', 'aid'))
         assert session.query(Album).count() == len(objects) + 1
         row = session.query(Album).filter_by(id=4).first()
         assert row.id == 4
         assert row.name == 'The Works'
+
+
+class TestConnections:
+
+    @staticmethod
+    def _verify_rollback(session, objects):
+        assert session.query(Album).count() == len(objects)
+        row = session.query(Album).filter_by(id=4).first()
+        assert row is None
+
+    @staticmethod
+    def _verify_commit(session, objects):
+        assert session.query(Album).count() == len(objects) + 1
+        row = session.query(Album).filter_by(id=4).first()
+        assert row.id == 4
+        assert row.name == 'The Works'
+
+    def _test_transactions(self, session, conn, sio, objects):
+        # Test rollback
+        sio.seek(0)
+        session.execute('begin;')
+        copy_from(sio, Album, conn)
+        session.execute('rollback;')
+        self._verify_rollback(session, objects)
+        # Test commit
+        sio.seek(0)
+        session.execute('begin;')
+        copy_from(sio, Album, conn)
+        session.execute('commit;')
+        self._verify_commit(session, objects)
+
+    def test_engine(self, session, objects):
+        sio = io.StringIO()
+        sio.write(u'\t'.join(['4', 'The Works']))
+        sio.seek(0)
+        copy_from(sio, Album, session.connection().engine)
+        self._verify_commit(session, objects)
+
+    def test_connection_with_txn(self, session, objects):
+        sio = io.StringIO()
+        sio.write(u'\t'.join(['4', 'The Works']))
+        conn = session.connection()
+        self._test_transactions(session, conn, sio, objects)
+
+    def test_raw_connection_with_txn(self, session, objects):
+        sio = io.StringIO()
+        sio.write(u'\t'.join(['4', 'The Works']))
+        raw_conn = session.connection().connection
+        self._test_transactions(session, raw_conn, sio, objects)