def test_poke_raise_exception_on_terminal_state(self, mock_hook):
mock_get_job = mock_hook.return_value.get_job
mock_fetch_job_autoscaling_events_by_id = mock_hook.return_value.fetch_job_autoscaling_events_by_id
callback = mock.MagicMock()
task = DataflowJobAutoScalingEventsSensor(
task_id=TEST_TASK_ID,
job_id=TEST_JOB_ID,
callback=callback,
fail_on_terminal_state=True,
location=TEST_LOCATION,
project_id=TEST_PROJECT_ID,
gcp_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
mock_get_job.return_value = {"id": TEST_JOB_ID, "currentState": DataflowJobStatus.JOB_STATE_DONE}
with self.assertRaisesRegex(
AirflowException,
f"Job with id '{TEST_JOB_ID}' is already in terminal state: "
f"{DataflowJobStatus.JOB_STATE_DONE}",
):
task.poke(mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
mock_fetch_job_autoscaling_events_by_id.assert_not_called()
callback.assert_not_called()
def test_poke(self, job_current_state, fail_on_terminal_state, mock_hook):
mock_get_job = mock_hook.return_value.get_job
mock_fetch_job_autoscaling_events_by_id = mock_hook.return_value.fetch_job_autoscaling_events_by_id
callback = mock.MagicMock()
task = DataflowJobAutoScalingEventsSensor(
task_id=TEST_TASK_ID,
job_id=TEST_JOB_ID,
callback=callback,
fail_on_terminal_state=fail_on_terminal_state,
location=TEST_LOCATION,
project_id=TEST_PROJECT_ID,
gcp_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
mock_get_job.return_value = {"id": TEST_JOB_ID, "currentState": job_current_state}
results = task.poke(mock.MagicMock())
self.assertEqual(callback.return_value, results)
mock_hook.assert_called_once_with(
gcp_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
mock_fetch_job_autoscaling_events_by_id.assert_called_once_with(
job_id=TEST_JOB_ID, project_id=TEST_PROJECT_ID, location=TEST_LOCATION
)
callback.assert_called_once_with(mock_fetch_job_autoscaling_events_by_id.return_value)