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
import os
import signal
import subprocess
from unittest import mock
from django.db.backends.postgresql.client import DatabaseClient
from django.test import SimpleTestCase
class PostgreSqlDbshellCommandTestCase(SimpleTestCase):
def _run_it(self, dbinfo):
"""
That function invokes the runshell command, while mocking
subprocess.run(). It returns a 2-tuple with:
- The command line list
- The dictionary of PG* environment variables, or {}.
"""
def _mock_subprocess_run(*args, env=os.environ, **kwargs):
self.subprocess_args = list(*args)
# PostgreSQL environment variables.
self.pg_env = {key: env[key] for key in env if key.startswith('PG')}
return subprocess.CompletedProcess(self.subprocess_args, 0)
with mock.patch('subprocess.run', new=_mock_subprocess_run):
DatabaseClient.runshell_db(dbinfo)
return self.subprocess_args, self.pg_env
def test_basic(self):
self.assertEqual(
self._run_it({
'database': 'dbname',
'user': 'someuser',
'password': 'somepassword',
'host': 'somehost',
'port': '444',
}), (
['psql', '-U', 'someuser', '-h', 'somehost', '-p', '444', 'dbname'],
{'PGPASSWORD': 'somepassword'},
)
)
def test_nopass(self):
self.assertEqual(
self._run_it({
'database': 'dbname',
'user': 'someuser',
'host': 'somehost',
'port': '444',
}), (
['psql', '-U', 'someuser', '-h', 'somehost', '-p', '444', 'dbname'],
{},
)
)
def test_ssl_certificate(self):
self.assertEqual(
self._run_it({
'database': 'dbname',
'user': 'someuser',
'host': 'somehost',
'port': '444',
'sslmode': 'verify-ca',
'sslrootcert': 'root.crt',
'sslcert': 'client.crt',
'sslkey': 'client.key',
}), (
['psql', '-U', 'someuser', '-h', 'somehost', '-p', '444', 'dbname'],
{
'PGSSLCERT': 'client.crt',
'PGSSLKEY': 'client.key',
'PGSSLMODE': 'verify-ca',
'PGSSLROOTCERT': 'root.crt',
},
)
)
def test_column(self):
self.assertEqual(
self._run_it({
'database': 'dbname',
'user': 'some:user',
'password': 'some:password',
'host': '::1',
'port': '444',
}), (
['psql', '-U', 'some:user', '-h', '::1', '-p', '444', 'dbname'],
{'PGPASSWORD': 'some:password'},
)
)
def test_accent(self):
username = 'rôle'
password = 'sésame'
self.assertEqual(
self._run_it({
'database': 'dbname',
'user': username,
'password': password,
'host': 'somehost',
'port': '444',
}), (
['psql', '-U', username, '-h', 'somehost', '-p', '444', 'dbname'],
{'PGPASSWORD': password},
)
)
def test_sigint_handler(self):
"""SIGINT is ignored in Python and passed to psql to abort quries."""
def _mock_subprocess_run(*args, **kwargs):
handler = signal.getsignal(signal.SIGINT)
self.assertEqual(handler, signal.SIG_IGN)
sigint_handler = signal.getsignal(signal.SIGINT)
# The default handler isn't SIG_IGN.
self.assertNotEqual(sigint_handler, signal.SIG_IGN)
with mock.patch('subprocess.run', new=_mock_subprocess_run):
DatabaseClient.runshell_db({})
# dbshell restores the original handler.
self.assertEqual(sigint_handler, signal.getsignal(signal.SIGINT))