summaryrefslogtreecommitdiff
path: root/django/test/utils.py
blob: bde323fa4e4ec6a85b554957fc47d7810ba18b32 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
import sys, time
from django.conf import settings
from django.db import backend, connect, connection, connection_info, connections

# The prefix to put on the default database name when creating
# the test database.
TEST_DATABASE_PREFIX = 'test_'

def _set_autocommit(connection):
    "Make sure a connection is in autocommit mode."
    if hasattr(connection.connection, "autocommit"):
        connection.connection.autocommit(True)
    elif hasattr(connection.connection, "set_isolation_level"):
        connection.connection.set_isolation_level(0)
        
def create_test_db(verbosity=1, autoclobber=False):
    if verbosity >= 1:
        print "Creating test database..."
        
    # If we're using SQLite, it's more convenient to test against an
    # in-memory database.
    if settings.DATABASE_ENGINE == "sqlite3":
        TEST_DATABASE_NAME = ":memory:"
        if verbosity >= 2:
            print "Using in-memory sqlite database for testing"
    else:
        if settings.TEST_DATABASE_NAME:
            TEST_DATABASE_NAME = settings.TEST_DATABASE_NAME
        else:
            TEST_DATABASE_NAME = TEST_DATABASE_PREFIX + settings.DATABASE_NAME

        qn = backend.quote_name
        # Create the test database and connect to it. We need to autocommit
        # if the database supports it because PostgreSQL doesn't allow 
        # CREATE/DROP DATABASE statements within transactions.
        cursor = connection.cursor()
        _set_autocommit(connection)
        try:
            cursor.execute("CREATE DATABASE %s" % qn(db_name))
        except Exception, e:            
            sys.stderr.write("Got an error creating the test database: %s\n" % e)
            if not autoclobber:
                confirm = raw_input("It appears the test database, %s, already exists. Type 'yes' to delete it, or 'no' to cancel: " % db_name)
            if autoclobber or confirm == 'yes':
                try:
                    if verbosity >= 1:
                        print "Destroying old test database..."                
                    cursor.execute("DROP DATABASE %s" % qn(db_name))
                    if verbosity >= 1:
                        print "Creating test database..."
                    cursor.execute("CREATE DATABASE %s" % qn(db_name))
                except Exception, e:
                    sys.stderr.write("Got an error recreating the test database: %s\n" % e)
                    sys.exit(2)
            else:
                print "Tests cancelled."
                sys.exit(1)
        # Close the old connection
        connection.close()

        # Get a cursor (even though we don't need one yet). This has
        # the side effect of initializing the test database.
        cursor = connection.cursor()        

    # Fill OTHER_DATABASES with the TEST_DATABASES settings,
    # and connect each named connection to the test database, using
    # a separate connection instance for each (so, eg, transactions don't
    # collide)
    test_databases = {}
    for db_name in settings.TEST_DATABASES:
        if settings.DATABASE_ENGINE == 'sqlite3':
            full_name = TEST_DATABASE_NAME
        else:
            full_name = TEST_DATABASE_NAME + db_name
        db_st = {'DATABASE_NAME': full_name}
        if db_name in settings.TEST_DATABASE_MODELS:
            db_st['MODELS'] = settings.TEST_DATABASE_MODELS.get(db_name, [])
        test_databases[db_name] = db_st
        connections[db_name] = connect(connection_info.settings)
        connections[db_name].connection.cursor() # Initialize it
    settings.OTHER_DATABASES = test_databases

def destroy_test_db(old_database_name, old_databases, verbosity=1):
    # Unless we're using SQLite, remove the test database to clean up after
    # ourselves. Connect to the previous database (not the test database)
    # to do so, because it's not allowed to delete a database while being
    # connected to it.
    if verbosity >= 1:
        print "Destroying test database..."
    if settings.DATABASE_ENGINE != "sqlite3": 
        connection.close()
        TEST_DATABASE_NAME = settings.DATABASE_NAME
        settings.DATABASE_NAME = old_database_name
        settings.OTHER_DATABASES = old_databases
        cursor = connection.cursor()
        _set_autocommit(connection)
        time.sleep(1) # To avoid "database is being accessed by other users" errors.
        cursor.execute("DROP DATABASE %s" % backend.quote_name(TEST_DATABASE_NAME))
        connection.close()