def test_get_unknown_workflow_status_case1(self, MockPost, MockBoto3):
service = MwaaWorkflowService("us-west-2", "access_key", "access_data")
service.client.create_cli_token = MagicMock(
return_value={
"CliToken": "cli_token",
"WebServerHostname": "host_name"
})
service.parse_response_json_result = MagicMock(return_value=[
{
"dag_id": "test_dag",
"task_id": "task1",
"state": "success"
},
{
"dag_id": "test_dag",
"task_id": "task2"
},
])
status = service.get_workflow_status(
{
"env_name": "test_env",
"dag_id": "test_dag"
}, "test_run_id")
service.client.create_cli_token.assert_called_with(Name="test_env")
self.assertEqual(status, WorkflowStatus.UNKNOWN)
def test_start_workflow_without_conf(self, MockPost, MockBoto3):
service = MwaaWorkflowService("us-west-2", "access_key", "access_data")
service.client.create_cli_token = MagicMock(
return_value={
"CliToken": "cli_token",
"WebServerHostname": "host_name"
})
service.parse_response_plain_result = MagicMock(
return_value=("test_run_id, externally triggered: True",
"test stderr"))
service.start_workflow({
"env_name": "test_env",
"dag_id": "test_dag"
}, "test_run_id")
service.client.create_cli_token.assert_called_with(Name="test_env")
def test_start_workflow_with_invalid_dag_id(self, MockPost, MockBoto3):
service = MwaaWorkflowService("us-west-2", "access_key", "access_data")
service.client.create_cli_token = MagicMock(
return_value={
"CliToken": "cli_token",
"WebServerHostname": "host_name"
})
service.parse_response_plain_result = MagicMock(
return_value=("test stdout",
"Dag id test_dag not found in DagModel"))
with self.assertRaises(Exception):
service.start_workflow(
{
"env_name": "test_env",
"dag_id": "test_dag"
},
"test_run_id",
{"conf1": "test_conf"},
)