diff --git a/.gitignore b/.gitignore index 21d7495..eaf7061 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ +.env *.egg *.egg-info/ *.pyc diff --git a/postdoc.py b/postdoc.py index d4ac9b4..11a2947 100755 --- a/postdoc.py +++ b/postdoc.py @@ -56,9 +56,6 @@ def connect_bits(meta): def pg_command(command, meta): """Construct the command.""" bits = [] - # password as environment varariable - if meta.password: - bits.append('PGPASSWORD=%s' % meta.password) # command to run bits.append(command) # connection params @@ -72,22 +69,24 @@ def pg_command(command, meta): def main(): if len(sys.argv) < 2: exit('Usage: phd COMMAND [additional-options]\n\n' - ' ERROR: Must give a COMMAND like psql, createdb, dropdb' - ) + ' ERROR: Must give a COMMAND like psql, createdb, dropdb') if sys.argv[1] not in VALID_COMMANDS: exit('Usage: phd COMMAND [additional-options]\n\n' - ' ERROR: "%s" is not a known postgres command' % sys.argv[1] - ) + ' ERROR: "%s" is not a known postgres command' % sys.argv[1]) try: - tokens = pg_command(sys.argv[1], get_uri()) + meta = get_uri() + tokens = pg_command(sys.argv[1], meta) except AttributeError: exit('Usage: phd COMMAND [additional-options]\n\n' - ' ERROR: DATABASE_URL is not set' - ) + ' ERROR: DATABASE_URL is not set') + env = os.environ.copy() + # password as environment varariable + if meta.password: + env['PGPASSWORD'] = meta.password # pass any other flags the user set along tokens.extend(sys.argv[2:]) - subprocess.call(tokens) + subprocess.call(tokens, env=env) # TODO test that PGPASS is in env if __name__ == '__main__': diff --git a/setup.py b/setup.py index d89c116..b3cd69d 100644 --- a/setup.py +++ b/setup.py @@ -3,7 +3,7 @@ setup( name='postdoc', - version='0.1.2', + version='0.1.3', description='A helper for Postgres + Docker that works for free', long_description=open('README.rst').read(), author='Chris Chang', @@ -16,7 +16,7 @@ ], }, license='Apache', - tests_require = [ + tests_require=[ 'mock==1.0.1', ], test_suite='test_postdoc', diff --git a/test_postdoc.py b/test_postdoc.py index 9eb7868..9ecbe61 100644 --- a/test_postdoc.py +++ b/test_postdoc.py @@ -1,5 +1,6 @@ # -*- coding: utf-8 -*- +import os try: # for python 2.6 compatibility import unittest2 as unittest @@ -42,14 +43,14 @@ def test_pg_command_assembles_bits_in_right_order(self): self.assertEqual(postdoc.pg_command('foo', meta), ['foo', 'lol', 'database']) - def test_pg_command_prepends_password(self): + def test_pg_command_ignores_password(self): meta = type('mock', (object, ), {'username': '', 'hostname': '', 'port': '', 'password': 'oops', 'path': '/database'}) with mock.patch('postdoc.connect_bits') as mock_bits: mock_bits.return_value = ['rofl'] self.assertEqual(postdoc.pg_command('bar', meta), - ['PGPASSWORD=oops', 'bar', 'rofl', 'database']) + ['bar', 'rofl', 'database']) def test_main_exits_with_no_command(self): with mock.patch('postdoc.sys') as mock_sys: @@ -69,7 +70,8 @@ def test_main_exits_with_missing_env(self): mock_sys = mock.MagicMock() mock_sys.argv = ['argv1', 'psql', 'argv3', 'argv4'] - with mock.patch.multiple(postdoc, + with mock.patch.multiple( + postdoc, subprocess=mock_subprocess, pg_command=mock_pg_command, sys=mock_sys, @@ -77,22 +79,51 @@ def test_main_exits_with_missing_env(self): with self.assertRaises(SystemExit): postdoc.main() + def test_main_passes_password_in_env(self): + my_password = 'oops' + meta = type('mock', (object, ), + {'password': my_password}) + self.assertNotIn('DATABASE_URL', os.environ, + msg="Re-run tests in an environment without DATABASE_URL") + mock_subprocess = mock.MagicMock() + mock_pg_command = mock.MagicMock(return_value=['pg_command']) + mock_get_uri = mock.MagicMock(return_value=meta) + mock_sys = mock.MagicMock() + mock_sys.argv = ['foo', 'psql'] + + with mock.patch.multiple( + postdoc, + subprocess=mock_subprocess, + pg_command=mock_pg_command, + get_uri=mock_get_uri, + sys=mock_sys, + ): + postdoc.main() + self.assertEqual( + mock_subprocess.call.call_args[1]['env']['PGPASSWORD'], + my_password) + def test_main_appends_additional_flags(self): + self.assertNotIn('DATABASE_URL', os.environ, + msg="Re-run tests in an environment without DATABASE_URL") mock_subprocess = mock.MagicMock() mock_pg_command = mock.MagicMock(return_value=['pg_command']) + mock_get_uri = mock.MagicMock() mock_sys = mock.MagicMock() mock_sys.argv = ['argv1', 'psql', 'argv3', 'argv4'] - mock_get_uri = mock.MagicMock() - with mock.patch.multiple(postdoc, + with mock.patch.multiple( + postdoc, subprocess=mock_subprocess, pg_command=mock_pg_command, get_uri=mock_get_uri, sys=mock_sys, ): postdoc.main() - mock_subprocess.call.assert_called_once_with( - ['pg_command', 'argv3', 'argv4']) + self.assertEqual( + mock_subprocess.call.call_args[0][0], + ['pg_command', 'argv3', 'argv4'] + ) if __name__ == '__main__':