def test_poke_success(self, mock_hook):
mock_hook.return_value.get_execution.return_value = mock.MagicMock(state=Execution.State.SUCCEEDED)
op = WorkflowExecutionSensor(
task_id="test_task",
workflow_id=WORKFLOW_ID,
execution_id=EXECUTION_ID,
location=LOCATION,
project_id=PROJECT_ID,
retry=RETRY,
request_timeout=TIMEOUT,
metadata=METADATA,
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
result = op.poke({})
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
mock_hook.return_value.get_execution.assert_called_once_with(
workflow_id=WORKFLOW_ID,
execution_id=EXECUTION_ID,
location=LOCATION,
project_id=PROJECT_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert result is True
def test_poke_failure(self, mock_hook):
mock_hook.return_value.get_execution.return_value = mock.MagicMock(state=Execution.State.FAILED)
op = WorkflowExecutionSensor(
task_id="test_task",
workflow_id=WORKFLOW_ID,
execution_id=EXECUTION_ID,
location=LOCATION,
project_id=PROJECT_ID,
retry=RETRY,
request_timeout=TIMEOUT,
metadata=METADATA,
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
with pytest.raises(AirflowException):
op.poke({})
create_execution = WorkflowsCreateExecutionOperator(
task_id="create_execution",
location=LOCATION,
project_id=PROJECT_ID,
execution=EXECUTION,
workflow_id=WORKFLOW_ID,
)
# [END how_to_create_execution]
create_execution_id = create_execution.output["execution_id"]
# [START how_to_wait_for_execution]
wait_for_execution = WorkflowExecutionSensor(
task_id="wait_for_execution",
location=LOCATION,
project_id=PROJECT_ID,
workflow_id=WORKFLOW_ID,
execution_id=create_execution_id,
)
# [END how_to_wait_for_execution]
# [START how_to_get_execution]
get_execution = WorkflowsGetExecutionOperator(
task_id="get_execution",
location=LOCATION,
project_id=PROJECT_ID,
workflow_id=WORKFLOW_ID,
execution_id=create_execution_id,
)
# [END how_to_get_execution]
# [START how_to_create_execution]
create_execution = WorkflowsCreateExecutionOperator(
task_id="create_execution",
location=LOCATION,
project_id=PROJECT_ID,
execution=EXECUTION,
workflow_id=WORKFLOW_ID,
)
# [END how_to_create_execution]
# [START how_to_wait_for_execution]
wait_for_execution = WorkflowExecutionSensor(
task_id="wait_for_execution",
location=LOCATION,
project_id=PROJECT_ID,
workflow_id=WORKFLOW_ID,
execution_id=
'{{ task_instance.xcom_pull("create_execution", key="execution_id") }}',
)
# [END how_to_wait_for_execution]
# [START how_to_get_execution]
get_execution = WorkflowsGetExecutionOperator(
task_id="get_execution",
location=LOCATION,
project_id=PROJECT_ID,
workflow_id=WORKFLOW_ID,
execution_id=
'{{ task_instance.xcom_pull("create_execution", key="execution_id") }}',
)
# [END how_to_get_execution]