← Back to branch summary

~james-w/ubuntu/lucid/psycopg2/precise-backport

« back to all changes in this revision

Viewing changes to tests/test_green.py

  • Committer: Mikhail Turov
  • Date: 2010-07-28 20:30:01 UTC
  • mfrom: (5.1.9 sid)
  • Revision ID: groldster@gmail.com-20100728203001-u1dv46tnd3s02ejg
Tags: 2.2.1-1ubuntu1
releasing version 2.2.1-1ubuntu1

Show diffs side-by-side

added added

removed removed

Lines of Context:
tests/test_green.py
 
1
#!/usr/bin/env python
 
2
 
 
3
import unittest
 
4
import psycopg2
 
5
import psycopg2.extensions
 
6
import psycopg2.extras
 
7
import tests
 
8
 
 
9
class ConnectionStub(object):
 
10
    """A `connection` wrapper allowing analysis of the `poll()` calls."""
 
11
    def __init__(self, conn):
 
12
        self.conn = conn
 
13
        self.polls = []
 
14
 
 
15
    def fileno(self):
 
16
        return self.conn.fileno()
 
17
 
 
18
    def poll(self):
 
19
        rv = self.conn.poll()
 
20
        self.polls.append(rv)
 
21
        return rv
 
22
 
 
23
class GreenTests(unittest.TestCase):
 
24
    def connect(self):
 
25
        return psycopg2.connect(tests.dsn)
 
26
 
 
27
    def setUp(self):
 
28
        self._cb = psycopg2.extensions.get_wait_callback()
 
29
        psycopg2.extensions.set_wait_callback(psycopg2.extras.wait_select)
 
30
 
 
31
    def tearDown(self):
 
32
        psycopg2.extensions.set_wait_callback(self._cb)
 
33
 
 
34
    def set_stub_wait_callback(self, conn):
 
35
        stub = ConnectionStub(conn)
 
36
        psycopg2.extensions.set_wait_callback(
 
37
            lambda conn: psycopg2.extras.wait_select(stub))
 
38
        return stub
 
39
 
 
40
    def test_flush_on_write(self):
 
41
        # a very large query requires a flush loop to be sent to the backend
 
42
        conn = self.connect()
 
43
        stub = self.set_stub_wait_callback(conn)
 
44
        curs = conn.cursor()
 
45
        for mb in 1, 5, 10, 20, 50:
 
46
            size = mb * 1024 * 1024
 
47
            del stub.polls[:]
 
48
            curs.execute("select %s;", ('x' * size,))
 
49
            self.assertEqual(size, len(curs.fetchone()[0]))
 
50
            if stub.polls.count(psycopg2.extensions.POLL_WRITE) > 1:
 
51
                return
 
52
 
 
53
        self.fail("sending a large query didn't trigger block on write.")
 
54
 
 
55
    def test_error_in_callback(self):
 
56
        conn = self.connect()
 
57
        curs = conn.cursor()
 
58
        curs.execute("select 1")  # have a BEGIN
 
59
        curs.fetchone()
 
60
 
 
61
        # now try to do something that will fail in the callback
 
62
        psycopg2.extensions.set_wait_callback(lambda conn: 1/0)
 
63
        self.assertRaises(ZeroDivisionError, curs.execute, "select 2")
 
64
 
 
65
        # check that the connection is left in an usable state
 
66
        psycopg2.extensions.set_wait_callback(psycopg2.extras.wait_select)
 
67
        conn.rollback()
 
68
        curs.execute("select 2")
 
69
        self.assertEqual(2, curs.fetchone()[0])
 
70
 
 
71
 
 
72
def test_suite():
 
73
    return unittest.TestLoader().loadTestsFromName(__name__)
 
74
 
 
75
if __name__ == "__main__":
 
76
    unittest.main()