def test_poke_raise_exception(self, mock_hook):
mock_get_job = mock_hook.return_value.get_job
task = DataflowJobStatusSensor(
task_id=TEST_TASK_ID,
job_id=TEST_JOB_ID,
expected_statuses=DataflowJobStatus.JOB_STATE_RUNNING,
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_CANCELLED}
with self.assertRaisesRegex(
AirflowException,
f"Job with id '{TEST_JOB_ID}' is already in terminal state: "
f"{DataflowJobStatus.JOB_STATE_CANCELLED}",
):
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_get_job.assert_called_once_with(
job_id=TEST_JOB_ID, project_id=TEST_PROJECT_ID, location=TEST_LOCATION
)
def test_poke(self, expected_status, current_status, sensor_return, mock_hook):
mock_get_job = mock_hook.return_value.get_job
task = DataflowJobStatusSensor(
task_id=TEST_TASK_ID,
job_id=TEST_JOB_ID,
expected_statuses=expected_status,
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": current_status}
results = task.poke(mock.MagicMock())
self.assertEqual(sensor_return, 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_get_job.assert_called_once_with(
job_id=TEST_JOB_ID, project_id=TEST_PROJECT_ID, location=TEST_LOCATION
)
},
py_requirements=['apache-beam[gcp]==2.25.0'],
py_interpreter='python3',
py_system_site_packages=False,
dataflow_config={
"location": 'europe-west3',
"wait_until_finished": False
},
)
# [END howto_operator_start_python_job_async]
# [START howto_sensor_wait_for_job_status]
wait_for_python_job_async_done = DataflowJobStatusSensor(
task_id="wait-for-python-job-async-done",
job_id=
"{{task_instance.xcom_pull('start-python-job-async')['job_id']}}",
expected_statuses={DataflowJobStatus.JOB_STATE_DONE},
location='europe-west3',
)
# [END howto_sensor_wait_for_job_status]
# [START howto_sensor_wait_for_job_metric]
def check_metric_scalar_gte(metric_name: str, value: int) -> Callable:
"""Check is metric greater than equals to given value."""
def callback(metrics: List[Dict]) -> bool:
dag_native_python_async.log.info("Looking for '%s' >= %d",
metric_name, value)
for metric in metrics:
context = metric.get("name", {}).get("context", {})
original_name = context.get("original_name", "")
runner="DataflowRunner",
py_file=GCS_PYTHON_DATAFLOW_ASYNC,
pipeline_options={
'tempLocation': GCS_TMP,
'stagingLocation': GCS_STAGING,
'output': GCS_OUTPUT,
},
py_options=[],
py_requirements=['apache-beam[gcp]==2.26.0'],
py_interpreter='python3',
py_system_site_packages=False,
dataflow_config=DataflowConfiguration(
job_name='{{task.task_id}}',
project_id=GCP_PROJECT_ID,
location="us-central1",
wait_until_finished=False,
),
)
wait_for_python_job_dataflow_runner_async_done = DataflowJobStatusSensor(
task_id="wait-for-python-job-async-done",
job_id=
"{{task_instance.xcom_pull('start_python_job_dataflow_runner_async')['dataflow_job_id']}}",
expected_statuses={DataflowJobStatus.JOB_STATE_DONE},
project_id=GCP_PROJECT_ID,
location='us-central1',
)
start_python_job_dataflow_runner_async >> wait_for_python_job_dataflow_runner_async_done
# [END howto_operator_start_python_dataflow_runner_pipeline_async_gcs_file]
