← Back to branch summary

~ubuntuone-pqm-team/canonical-identity-provider/trunk

« back to all changes in this revision

Viewing changes to identityprovider/tests/mockdb.py

  • Committer: Danny Tamez
  • Date: 2010-04-21 15:29:24 UTC
  • Revision ID: danny.tamez@canonical.com-20100421152924-lq1m92tstk2iz75a
Canonical SSO Provider (Open Source) - Initial Commit

Show diffs side-by-side

added added

removed removed

Lines of Context:
identityprovider/tests/mockdb.py
 
1
# Copyright 2010 Canonical Ltd.  This software is licensed under the
 
2
# GNU Affero General Public License version 3 (see the file LICENSE).
 
3
 
 
4
import psycopg2
 
5
from psycopg2 import DatabaseError
 
6
 
 
7
 
 
8
class MockConnection(object):
 
9
    def __init__(self, connection, psycopg):
 
10
        self.connection = connection
 
11
        self.psycopg = psycopg
 
12
 
 
13
    def cursor(self):
 
14
        if self.psycopg.fail_next_N_connections > 0:
 
15
            self.psycopg.fail_next_N_connections -= 1
 
16
            raise DatabaseError('Mock error.')
 
17
        return self.connection.cursor()
 
18
 
 
19
    def __getattr__(self, attr):
 
20
        if attr in self.__dict__:
 
21
            return self.__dict__[attr]
 
22
        else:
 
23
            return getattr(self.connection, attr)
 
24
 
 
25
 
 
26
class MockPsycopg(object):
 
27
    def __init__(self):
 
28
        self.fail_next_N_connections = 0
 
29
        self.dsn = None
 
30
 
 
31
    def connect(self, *args, **kwargs):
 
32
        if self.fail_next_N_connections > 0:
 
33
            self.fail_next_N_connections -= 1
 
34
            raise DatabaseError('Mock error.')
 
35
        if self.dsn is not None:
 
36
            args = []
 
37
            kwargs = {'dsn': self.dsn}
 
38
        return MockConnection(psycopg2.connect(*args, **kwargs), self)
 
39
 
 
40
    def fixate_connection(self, dsn):
 
41
        """ Fixate the database connection information.
 
42
 
 
43
        This function ensures that further connections will be
 
44
        attempted always using the same DSN, so that we can play
 
45
        around with the DATABASE_* settings during tests and still
 
46
        keep connecting to our only test database. """
 
47
        self.dsn = dsn
 
48
 
 
49
    def release_connection(self):
 
50
        """ Release the database connection information.
 
51
 
 
52
        Forget about stored connection information, so that next
 
53
        call to connect() really connects to the database you ask
 
54
        for. """
 
55
        self.dsn = None
 
56
 
 
57
mock = MockPsycopg()
 
58
connect = mock.connect
 
59
fixate_connection = mock.fixate_connection
 
60
release_connection = mock.release_connection
 
61
 
 
62
 
 
63
def fail_next_N_connections(n):
 
64
    mock.fail_next_N_connections = n