1 files changed, 122 insertions, 14 deletions
diff --git a/test/engine/test_reconnect.py b/test/engine/test_reconnect.py
index 6b283654b..86f646f33 100644
--- a/test/engine/test_reconnect.py
+++ b/test/engine/test_reconnect.py
@@ -11,7 +11,10 @@ from sqlalchemy import exc
 from sqlalchemy.testing import fixtures
 from sqlalchemy.testing.engines import testing_engine
 
-class MockDisconnect(Exception):
+class MockError(Exception):
+    pass
+
+class MockDisconnect(MockError):
     pass
 
 class MockDBAPI(object):
@@ -20,17 +23,23 @@ class MockDBAPI(object):
         self.connections = weakref.WeakKeyDictionary()
     def connect(self, *args, **kwargs):
         return MockConnection(self)
-    def shutdown(self):
+    def shutdown(self, explode='execute'):
         for c in self.connections:
-            c.explode[0] = True
-    Error = MockDisconnect
+            c.explode = explode
+    Error = MockError
 
 class MockConnection(object):
     def __init__(self, dbapi):
         dbapi.connections[self] = True
-        self.explode = [False]
+        self.explode = ""
     def rollback(self):
-        pass
+        if self.explode == 'rollback':
+            raise MockDisconnect("Lost the DB connection on rollback")
+        if self.explode == 'rollback_no_disconnect':
+            raise MockError(
+                "something broke on rollback but we didn't lose the connection")
+        else:
+            return
     def commit(self):
         pass
     def cursor(self):
@@ -43,8 +52,14 @@ class MockCursor(object):
         self.explode = parent.explode
         self.description = ()
     def execute(self, *args, **kwargs):
-        if self.explode[0]:
-            raise MockDisconnect("Lost the DB connection")
+        if self.explode == 'execute':
+            raise MockDisconnect("Lost the DB connection on execute")
+        elif self.explode in ('execute_no_disconnect', ):
+            raise MockError(
+                "something broke on execute but we didn't lose the connection")
+        elif self.explode in ('rollback', 'rollback_no_disconnect'):
+            raise MockError(
+                "something broke on execute but we didn't lose the connection")
         else:
             return
     def close(self):
@@ -167,12 +182,10 @@ class MockReconnectTest(fixtures.TestBase):
 
         dbapi.shutdown()
 
-        # raises error
-        try:
-            conn.execute(select([1]))
-            assert False
-        except tsa.exc.DBAPIError:
-            pass
+        assert_raises(
+            tsa.exc.DBAPIError,
+            conn.execute, select([1])
+        )
 
         assert not conn.closed
         assert conn.invalidated
@@ -186,6 +199,101 @@ class MockReconnectTest(fixtures.TestBase):
         assert not conn.invalidated
         assert len(dbapi.connections) == 1
 
+    def test_invalidated_close(self):
+        conn = db.connect()
+
+        dbapi.shutdown()
+
+        assert_raises(
+            tsa.exc.DBAPIError,
+            conn.execute, select([1])
+        )
+
+        conn.close()
+        assert conn.closed
+        assert conn.invalidated
+        assert_raises_message(
+            tsa.exc.StatementError,
+            "This Connection is closed",
+            conn.execute, select([1])
+        )
+
+    def test_noreconnect_execute_plus_closewresult(self):
+        conn = db.connect(close_with_result=True)
+
+        dbapi.shutdown("execute_no_disconnect")
+
+        # raises error
+        assert_raises_message(
+            tsa.exc.DBAPIError,
+            "something broke on execute but we didn't lose the connection",
+            conn.execute, select([1])
+        )
+
+        assert conn.closed
+        assert not conn.invalidated
+
+    def test_noreconnect_rollback_plus_closewresult(self):
+        conn = db.connect(close_with_result=True)
+
+        dbapi.shutdown("rollback_no_disconnect")
+
+        # raises error
+        assert_raises_message(
+            tsa.exc.DBAPIError,
+            "something broke on rollback but we didn't lose the connection",
+            conn.execute, select([1])
+        )
+
+        assert conn.closed
+        assert not conn.invalidated
+
+        assert_raises_message(
+            tsa.exc.StatementError,
+            "This Connection is closed",
+            conn.execute, select([1])
+        )
+
+    def test_reconnect_on_reentrant(self):
+        conn = db.connect()
+
+        conn.execute(select([1]))
+
+        assert len(dbapi.connections) == 1
+
+        dbapi.shutdown("rollback")
+
+        # raises error
+        assert_raises_message(
+            tsa.exc.DBAPIError,
+            "Lost the DB connection on rollback",
+            conn.execute, select([1])
+        )
+
+        assert not conn.closed
+        assert conn.invalidated
+
+    def test_reconnect_on_reentrant_plus_closewresult(self):
+        conn = db.connect(close_with_result=True)
+
+        dbapi.shutdown("rollback")
+
+        # raises error
+        assert_raises_message(
+            tsa.exc.DBAPIError,
+            "Lost the DB connection on rollback",
+            conn.execute, select([1])
+        )
+
+        assert conn.closed
+        assert conn.invalidated
+
+        assert_raises_message(
+            tsa.exc.StatementError,
+            "This Connection is closed",
+            conn.execute, select([1])
+        )
+
 class CursorErrTest(fixtures.TestBase):
 
     def setup(self):