Merge remote branch 'piro/fix22' into python2

1 files changed, 45 insertions, 0 deletions

diff --git a/tests/test_async.py b/tests/test_async.py
index 5aad8d3..da084bd 100644..100755
--- a/tests/test_async.py
+++ b/tests/test_async.py
@@ -14,6 +14,21 @@ else:
     import py3tests as tests
 
 
+class PollableStub(object):
+    """A 'pollable' wrapper allowing analysis of the `poll()` calls."""
+    def __init__(self, pollable):
+        self.pollable = pollable
+        self.polls = []
+
+    def fileno(self):
+        return self.pollable.fileno()
+
+    def poll(self):
+        rv = self.pollable.poll()
+        self.polls.append(rv)
+        return rv
+
+
 class AsyncTests(unittest.TestCase):
 
     def setUp(self):
@@ -274,5 +289,35 @@ class AsyncTests(unittest.TestCase):
         # it should be the result of the second query
         self.assertEquals(cur.fetchone()[0], "b" * 10000)
 
+    def test_async_subclass(self):
+        class MyConn(psycopg2.extensions.connection):
+            def __init__(self, dsn, async=0):
+                psycopg2.extensions.connection.__init__(self, dsn, async=async)
+
+        conn = psycopg2.connect(tests.dsn, connection_factory=MyConn, async=True)
+        self.assert_(isinstance(conn, MyConn))
+        self.assert_(not conn.issync())
+        conn.close()
+
+
+    def test_flush_on_write(self):
+        # a very large query requires a flush loop to be sent to the backend
+        curs = self.conn.cursor()
+        for mb in 1, 5, 10, 20, 50:
+            size = mb * 1024 * 1024
+            print "\nplease wait: sending", mb, "MB query to the server",
+            stub = PollableStub(curs)
+            curs.execute("select %s;", ('x' * size,))
+            self.wait(stub)
+            self.assertEqual(size, len(curs.fetchone()[0]))
+            if stub.polls.count(psycopg2.extensions.POLL_WRITE) > 1:
+                return
+
+        self.fail("sending a large query didn't trigger block on write.")
+
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)
+
+if __name__ == "__main__":
+    unittest.main()
+