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 the value of the PGPASSWORD environment variable, or None. """ def _mock_subprocess_run(*args, env=os.environ, **kwargs): self.subprocess_args = list(*args) self.pgpassword = env.get('PGPASSWORD') 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.pgpassword 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'], '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'], None, ) ) 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'], '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'], 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))