1 files changed, 90 insertions, 0 deletions
diff --git a/test/engine/test_pool.py b/test/engine/test_pool.py
index 10f490b48..2e4c2dc48 100644
--- a/test/engine/test_pool.py
+++ b/test/engine/test_pool.py
@@ -1328,6 +1328,96 @@ class QueuePoolTest(PoolTestBase):
         c2 = p.connect()
         assert c2.connection is not None
 
+class ResetOnReturnTest(PoolTestBase):
+    def _fixture(self, **kw):
+        dbapi = Mock()
+        return dbapi, pool.QueuePool(creator=lambda: dbapi.connect('foo.db'), **kw)
+
+    def test_plain_rollback(self):
+        dbapi, p = self._fixture(reset_on_return='rollback')
+
+        c1 = p.connect()
+        c1.close()
+        assert dbapi.connect().rollback.called
+        assert not dbapi.connect().commit.called
+
+    def test_plain_commit(self):
+        dbapi, p = self._fixture(reset_on_return='commit')
+
+        c1 = p.connect()
+        c1.close()
+        assert not dbapi.connect().rollback.called
+        assert dbapi.connect().commit.called
+
+    def test_plain_none(self):
+        dbapi, p = self._fixture(reset_on_return=None)
+
+        c1 = p.connect()
+        c1.close()
+        assert not dbapi.connect().rollback.called
+        assert not dbapi.connect().commit.called
+
+    def test_agent_rollback(self):
+        dbapi, p = self._fixture(reset_on_return='rollback')
+
+        class Agent(object):
+            def __init__(self, conn):
+                self.conn = conn
+
+            def rollback(self):
+                self.conn.special_rollback()
+
+            def commit(self):
+                self.conn.special_commit()
+
+        c1 = p.connect()
+        c1._reset_agent = Agent(c1)
+        c1.close()
+
+        assert dbapi.connect().special_rollback.called
+        assert not dbapi.connect().special_commit.called
+
+        assert not dbapi.connect().rollback.called
+        assert not dbapi.connect().commit.called
+
+        c1 = p.connect()
+        c1.close()
+        eq_(dbapi.connect().special_rollback.call_count, 1)
+        eq_(dbapi.connect().special_commit.call_count, 0)
+
+        assert dbapi.connect().rollback.called
+        assert not dbapi.connect().commit.called
+
+    def test_agent_commit(self):
+        dbapi, p = self._fixture(reset_on_return='commit')
+
+        class Agent(object):
+            def __init__(self, conn):
+                self.conn = conn
+
+            def rollback(self):
+                self.conn.special_rollback()
+
+            def commit(self):
+                self.conn.special_commit()
+
+        c1 = p.connect()
+        c1._reset_agent = Agent(c1)
+        c1.close()
+        assert not dbapi.connect().special_rollback.called
+        assert dbapi.connect().special_commit.called
+
+        assert not dbapi.connect().rollback.called
+        assert not dbapi.connect().commit.called
+
+        c1 = p.connect()
+        c1.close()
+
+        eq_(dbapi.connect().special_rollback.call_count, 0)
+        eq_(dbapi.connect().special_commit.call_count, 1)
+        assert not dbapi.connect().rollback.called
+        assert dbapi.connect().commit.called
+
 class SingletonThreadPoolTest(PoolTestBase):
 
     @testing.requires.threading_with_mock