def test_cli_shell_mysql(self, mock_tmp_file, mock_execute_interactive):
mock_tmp_file.return_value.__enter__.return_value.name = "/tmp/name"
db_command.shell(self.parser.parse_args(['db', 'shell']))
mock_execute_interactive.assert_called_once_with(['mysql', '--defaults-extra-file=/tmp/name'])
mock_tmp_file.return_value.__enter__.return_value.write.assert_called_once_with(
b'[client]\nhost = mysql\nuser = root\npassword = \nport = ' b'\ndatabase = airflow'
)
def test_cli_shell_mysql(self, mock_tmp_file, mock_subprocess):
db_command.shell(self.parser.parse_args(['db', 'shell']))
mock_subprocess.Popen.assert_called_once_with(
['mysql', '--defaults-extra-file=/tmp/name']
)
mock_tmp_file.return_value.__enter__.return_value.write.assert_called_once_with(
b'[client]\nhost = mysql\nuser = root\npassword = None\nport = None'
b'\ndatabase = airflow'
)
def test_cli_shell_postgres(self, mock_execute_interactive):
db_command.shell(self.parser.parse_args(['db', 'shell']))
mock_execute_interactive.assert_called_once_with(['psql'], env=mock.ANY)
_, kwargs = mock_execute_interactive.call_args
env = kwargs['env']
postgres_env = {k: v for k, v in env.items() if k.startswith('PG')}
assert {
'PGDATABASE': 'airflow',
'PGHOST': 'postgres',
'PGPASSWORD': '******',
'PGPORT': '5432',
'PGUSER': '******',
} == postgres_env
def test_cli_shell_postgres(self, mock_subprocess):
db_command.shell(self.parser.parse_args(['db', 'shell']))
mock_subprocess.Popen.assert_called_once_with(['psql'], env=mock.ANY)
_, kwargs = mock_subprocess.Popen.call_args
env = kwargs['env']
postgres_env = {k: v for k, v in env.items() if k.startswith('PG')}
self.assertEqual(
{
'PGDATABASE': 'airflow',
'PGHOST': 'postgres',
'PGPASSWORD': '******',
'PGPORT': '',
'PGUSER': '******'
}, postgres_env)
def test_cli_shell_sqlite(self, mock_execute_interactive):
db_command.shell(self.parser.parse_args(['db', 'shell']))
mock_execute_interactive.assert_called_once_with(
['sqlite3', '/root/airflow/airflow.db']
)
def test_cli_shell_invalid(self):
with self.assertRaisesRegex(AirflowException, r"Unknown driver: invalid\+psycopg2"):
db_command.shell(self.parser.parse_args(['db', 'shell']))
def test_cli_shell_sqlite(self, mock_subprocess):
db_command.shell(self.parser.parse_args(['db', 'shell']))
mock_subprocess.Popen.assert_called_once_with(
['sqlite3', '/root/airflow/airflow.db']
)
def test_cli_shell_invalid(self):
with pytest.raises(AirflowException, match=r"Unknown driver: invalid\+psycopg2"):
db_command.shell(self.parser.parse_args(['db', 'shell']))
