-
Notifications
You must be signed in to change notification settings - Fork 0
/
django_postgres_testing.py
151 lines (124 loc) · 4.76 KB
/
django_postgres_testing.py
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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
import copy
import shutil
import signal
import subprocess
import os
import tempfile
import psycopg2
import os.path
from glob import glob
from contextlib import closing
from django.db import connections
from django.test.runner import DiscoverRunner
SEARCH_PATHS = (['/usr/local/pgsql', '/usr/local'] +
glob('/usr/lib/postgresql/*') + # for Debian/Ubuntu
glob('/opt/local/lib/postgresql*')) # for MacPorts
def get_open_port():
import socket
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.bind(('', 0))
s.listen(1)
port = s.getsockname()[1]
s.close()
return port
class TemporaryPostgresRunner(DiscoverRunner):
def setup_databases(self, **kwargs):
postgres_binary_path = os.environ.get('POSTGRES_BIN')
if not postgres_binary_path:
# Try to find the initdb binary
initdb_path = shutil.which('initdb')
if initdb_path:
postgres_binary_path = os.path.dirname(initdb_path)
else:
for base_dir in SEARCH_PATHS:
path = os.path.join(base_dir, 'bin', 'initdb')
if os.path.exists(path):
postgres_binary_path = os.path.join(base_dir, 'bin')
break
self.postgres_socket_directory = tempfile.mkdtemp(prefix='pgsocket')
self.postgres_directory = tempfile.mkdtemp(prefix='pgdata')
self.postgres_port = get_open_port()
print('Initializing test Postgresql cluster...')
initdb_args = [
os.path.join(postgres_binary_path, 'initdb'),
'-A', 'trust',
'-U', 'postgres',
'--nosync',
'-D', self.postgres_directory
]
# Create the database...
init_process = subprocess.Popen(
initdb_args,
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
)
init_process.wait()
if init_process.returncode != 0:
raise Exception('Couldn\'t initialize database')
print('Starting test Postgresql cluster...')
postgres_args = [
os.path.join(postgres_binary_path, 'postgres'),
'-h 127.0.0.1',
'-F',
'-p', str(self.postgres_port),
'-D', self.postgres_directory,
'-k', self.postgres_socket_directory
]
self.postgres_process = subprocess.Popen(
postgres_args,
bufsize=1,
universal_newlines=1,
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
)
self._old_databases = copy.copy(connections.databases)
for name, settings in self._old_databases.items():
connections.databases[name]['HOST'] = 'localhost'
connections.databases[name]['PORT'] = self.postgres_port
# Let postgres start up...
while True:
try:
with closing(
psycopg2.connect(
host='localhost',
port=self.postgres_port,
user='postgres'
)
):
pass
except psycopg2.OperationalError:
# if not self.postgres_process.poll():
# raise Exception('Couldn\'t start postgres')
continue
else:
break
# Now let's create a user:
for name, settings in self._old_databases.items():
username = connections.databases[name]['USER']
createuser_args = [
os.path.join(postgres_binary_path, 'createuser'),
'-h', 'localhost',
'-p', str(self.postgres_port),
'-U', 'postgres',
'--createdb',
username
]
create_process = subprocess.Popen(
createuser_args,
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
)
create_process.wait()
if create_process.returncode != 0:
self.postgres_process.send_signal(signal.SIGINT)
self.postgres_process.wait()
shutil.rmtree(self.postgres_directory)
shutil.rmtree(self.postgres_socket_directory)
raise Exception('Couldn\'t create user')
return super(TemporaryPostgresRunner, self).setup_databases(**kwargs)
def teardown_databases(self, old_config, **kwargs):
ret = super(TemporaryPostgresRunner, self).teardown_databases(
old_config, **kwargs
)
self.postgres_process.send_signal(signal.SIGINT)
self.postgres_process.wait()
shutil.rmtree(self.postgres_directory)
shutil.rmtree(self.postgres_socket_directory)
return ret