def test_cli_connections_export_should_export_as_env_for_uppercase_file_extension( self, mock_file_open, mock_splittext ): output_filepath = '/tmp/connections.ENV' mock_splittext.return_value = (None, '.ENV') args = self.parser.parse_args( [ "connections", "export", output_filepath, ] ) connection_command.connections_export(args) expected_connections = [ "airflow_db=mysql://root:plainpassword@mysql/airflow\n" "druid_broker_default=druid://druid-broker:8082?endpoint=druid%2Fv2%2Fsql\n", "druid_broker_default=druid://druid-broker:8082?endpoint=druid%2Fv2%2Fsql\n" "airflow_db=mysql://root:plainpassword@mysql/airflow\n", ] mock_splittext.assert_called_once() mock_file_open.assert_called_once_with(output_filepath, 'w', -1, 'UTF-8', None) mock_file_open.return_value.write.assert_called_once_with(mock.ANY) self.assertIn(mock_file_open.return_value.write.call_args_list[0][0][0], expected_connections)
def test_cli_connections_export_should_export_as_yaml( self, mock_file_open, mock_splittext): output_filepath = '/tmp/connections.yaml' mock_splittext.return_value = (None, '.yaml') args = self.parser.parse_args([ "connections", "export", output_filepath, ]) connection_command.connections_export(args) expected_connections = ( "airflow_db:\n" " conn_type: mysql\n" " extra: null\n" " host: mysql\n" " login: root\n" " password: plainpassword\n" " port: null\n" " schema: airflow\n" "druid_broker_default:\n" " conn_type: druid\n" " extra: \'{\"endpoint\": \"druid/v2/sql\"}\'\n" " host: druid-broker\n" " login: null\n" " password: null\n" " port: 8082\n" " schema: null\n") mock_splittext.assert_called_once() mock_file_open.assert_called_once_with(output_filepath, 'w', -1, 'UTF-8', None) mock_file_open.return_value.write.assert_called_once_with( expected_connections)
def test_cli_connections_export_should_return_error_if_fetching_connections_fails( self, mock_session, mock_file_open, mock_splittext ): output_filepath = '/tmp/connections.json' def my_side_effect(_): raise Exception("dummy exception") mock_session.return_value.__enter__.return_value.query.return_value.order_by.side_effect = ( my_side_effect ) mock_splittext.return_value = (None, '.json') args = self.parser.parse_args( [ "connections", "export", output_filepath, ] ) with self.assertRaisesRegex(Exception, r"dummy exception"): connection_command.connections_export(args) mock_splittext.assert_called_once() mock_file_open.assert_called_once_with(output_filepath, 'w', -1, 'UTF-8', None) mock_file_open.return_value.write.assert_not_called()
def test_cli_connections_export_should_export_as_json(self, mock_file_open, mock_splittext): output_filepath = '/tmp/connections.json' mock_splittext.return_value = (None, '.json') args = self.parser.parse_args([ "connections", "export", output_filepath, ]) connection_command.connections_export(args) expected_connections = json.dumps({ "airflow_db": { "conn_type": "mysql", "host": "mysql", "login": "******", "password": "******", "schema": "airflow", "port": None, "extra": None, }, "druid_broker_default": { "conn_type": "druid", "host": "druid-broker", "login": None, "password": None, "schema": None, "port": 8082, "extra": "{\"endpoint\": \"druid/v2/sql\"}", } }, indent=2) mock_splittext.assert_called_once() mock_file_open.assert_called_once_with(output_filepath, 'w', -1, 'UTF-8', None) mock_file_open.return_value.write.assert_called_once_with(expected_connections)
def test_cli_connections_export_should_not_return_error_if_connections_is_empty( self, mock_session, mock_file_open, mock_splittext): output_filepath = '/tmp/connections.json' mock_session.return_value.__enter__.return_value.query.return_value.all.return_value = [] mock_splittext.return_value = (None, '.json') args = self.parser.parse_args([ "connections", "export", output_filepath, ]) connection_command.connections_export(args) mock_splittext.assert_called_once() mock_file_open.assert_called_once_with(output_filepath, 'w', -1, 'UTF-8', None) mock_file_open.return_value.write.assert_called_once_with('{}')
def test_cli_connections_export_should_export_as_env(self, mock_file_open, mock_splittext): output_filepath = '/tmp/connections.env' mock_splittext.return_value = (None, '.env') args = self.parser.parse_args([ "connections", "export", output_filepath, ]) connection_command.connections_export(args) expected_connections = ( "airflow_db=mysql://root:plainpassword@mysql/airflow\n" "druid_broker_default=druid://druid-broker:8082?endpoint=druid%2Fv2%2Fsql\n") mock_splittext.assert_called_once() mock_file_open.assert_called_once_with(output_filepath, 'w', -1, 'UTF-8', None) mock_file_open.return_value.write.assert_called_once_with(expected_connections)
def test_cli_connections_export_should_return_error_for_invalid_export_format(self, mock_file_open, mock_splittext): output_filepath = '/tmp/connections.invalid' mock_splittext.return_value = (None, '.invalid') args = self.parser.parse_args([ "connections", "export", output_filepath, ]) with self.assertRaisesRegex( SystemExit, r"Unsupported file format. The file must have the extension .yaml, .json, .env" ): connection_command.connections_export(args) mock_splittext.assert_called_once() mock_file_open.assert_called_once_with(output_filepath, 'w', -1, 'UTF-8', None) mock_file_open.return_value.write.assert_not_called()
def test_cli_connections_export_should_return_error_if_create_session_fails( self, mock_session, mock_file_open, mock_splittext): output_filepath = '/tmp/connections.json' def my_side_effect(): raise Exception("dummy exception") mock_session.side_effect = my_side_effect mock_splittext.return_value = (None, '.json') args = self.parser.parse_args([ "connections", "export", output_filepath, ]) with pytest.raises(Exception, match=r"dummy exception"): connection_command.connections_export(args) mock_splittext.assert_not_called() mock_file_open.assert_called_once_with(output_filepath, 'w', -1, 'UTF-8', None) mock_file_open.return_value.write.assert_not_called()