a rudimental reconnect/pool auto-dispose test. not super-comprehensive but better

than nothing, will close #516

1 files changed, 95 insertions, 0 deletions

diff --git a/test/engine/reconnect.py b/test/engine/reconnect.py
new file mode 100644
index 000000000..defc878ab
--- /dev/null
+++ b/test/engine/reconnect.py
@@ -0,0 +1,95 @@
+import testbase
+from sqlalchemy import create_engine, exceptions
+import gc, weakref, sys
+
+class MockDisconnect(Exception):
+    pass
+
+class MockDBAPI(object):
+    def __init__(self):
+        self.paramstyle = 'named'
+        self.connections = weakref.WeakKeyDictionary()
+    def connect(self, *args, **kwargs):
+        return MockConnection(self)
+        
+class MockConnection(object):
+    def __init__(self, dbapi):
+        self.explode = False
+        dbapi.connections[self] = True
+    def rollback(self):
+        pass
+    def commit(self):
+        pass
+    def cursor(self):
+        return MockCursor(explode=self.explode)
+    def close(self):
+        pass
+            
+class MockCursor(object):
+    def __init__(self, explode):
+        self.explode = explode
+        self.description = None
+    def execute(self, *args, **kwargs):
+        if self.explode:
+            raise MockDisconnect("Lost the DB connection")
+        else:
+            return
+    def close(self):
+        pass
+        
+class ReconnectTest(testbase.PersistTest):
+    def test_reconnect(self):
+        """test that an 'is_disconnect' condition will invalidate the connection, and additionally
+        dispose the previous connection pool and recreate."""
+        
+        dbapi = MockDBAPI()
+        
+        # create engine using our current dburi
+        db = create_engine('postgres://foo:bar@localhost/test', module=dbapi)
+        
+        # monkeypatch disconnect checker
+        db.dialect.is_disconnect = lambda e: isinstance(e, MockDisconnect)
+        
+        pid = id(db.connection_provider._pool)
+        
+        # make a connection
+        conn = db.connect()
+        
+        # connection works
+        conn.execute("SELECT 1")
+        
+        # create a second connection within the pool, which we'll ensure also goes away
+        conn2 = db.connect()
+        conn2.close()
+
+        # two connections opened total now
+        assert len(dbapi.connections) == 2
+
+        # set it to fail
+        conn.connection.connection.explode = True
+        
+        try:
+            # execute should fail
+            conn.execute("SELECT 1")
+            assert False
+        except exceptions.SQLAlchemyError, e:
+            pass
+        
+        # assert was invalidated
+        assert conn.connection.connection is None
+        
+        # close shouldnt break
+        conn.close()
+
+        assert id(db.connection_provider._pool) != pid
+        
+        # ensure all connections closed (pool was recycled)
+        assert len(dbapi.connections) == 0
+        
+        conn =db.connect()
+        conn.execute("SELECT 1")
+        conn.close()
+        assert len(dbapi.connections) == 1
+        
+if __name__ == '__main__':
+    testbase.main()
\ No newline at end of file