def execute(self, context: 'Context'):
hook = WorkflowsHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
self.log.info("Creating execution")
execution = hook.create_execution(
workflow_id=self.workflow_id,
execution=self.execution,
location=self.location,
project_id=self.project_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
execution_id = execution.name.split("/")[-1]
self.xcom_push(context, key="execution_id", value=execution_id)
WorkflowsExecutionLink.persist(
context=context,
task_instance=self,
location_id=self.location,
workflow_id=self.workflow_id,
execution_id=execution_id,
project_id=self.project_id or hook.project_id,
)
return Execution.to_dict(execution)
class TestWorkflowsHook:
def setup_method(self, _):
with mock.patch(BASE_PATH.format("GoogleBaseHook.__init__"),
new=mock_init):
self.hook = WorkflowsHook(gcp_conn_id="test") # pylint: disable=attribute-defined-outside-init
@mock.patch(BASE_PATH.format("WorkflowsHook._get_credentials"))
@mock.patch(BASE_PATH.format("WorkflowsHook.client_info"),
new_callable=mock.PropertyMock)
@mock.patch(BASE_PATH.format("WorkflowsClient"))
def test_get_workflows_client(self, mock_client, mock_client_info,
mock_get_credentials):
self.hook.get_workflows_client()
mock_client.assert_called_once_with(
credentials=mock_get_credentials.return_value,
client_info=mock_client_info.return_value,
)
@mock.patch(BASE_PATH.format("WorkflowsHook._get_credentials"))
@mock.patch(BASE_PATH.format("WorkflowsHook.client_info"),
new_callable=mock.PropertyMock)
@mock.patch(BASE_PATH.format("ExecutionsClient"))
def test_get_executions_client(self, mock_client, mock_client_info,
mock_get_credentials):
self.hook.get_executions_client()
mock_client.assert_called_once_with(
credentials=mock_get_credentials.return_value,
client_info=mock_client_info.return_value,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_workflows_client"))
def test_create_workflow(self, mock_client):
result = self.hook.create_workflow(
workflow=WORKFLOW,
workflow_id=WORKFLOW_ID,
location=LOCATION,
project_id=PROJECT_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.create_workflow.return_value == result
mock_client.return_value.create_workflow.assert_called_once_with(
request=dict(workflow=WORKFLOW,
workflow_id=WORKFLOW_ID,
parent=WORKFLOW_PARENT),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_workflows_client"))
def test_get_workflow(self, mock_client):
result = self.hook.get_workflow(
workflow_id=WORKFLOW_ID,
location=LOCATION,
project_id=PROJECT_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.get_workflow.return_value == result
mock_client.return_value.get_workflow.assert_called_once_with(
request=dict(name=WORKFLOW_NAME),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_workflows_client"))
def test_update_workflow(self, mock_client):
result = self.hook.update_workflow(
workflow=WORKFLOW,
update_mask=UPDATE_MASK,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.update_workflow.return_value == result
mock_client.return_value.update_workflow.assert_called_once_with(
request=dict(
workflow=WORKFLOW,
update_mask=UPDATE_MASK,
),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_workflows_client"))
def test_delete_workflow(self, mock_client):
result = self.hook.delete_workflow(
workflow_id=WORKFLOW_ID,
location=LOCATION,
project_id=PROJECT_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.delete_workflow.return_value == result
mock_client.return_value.delete_workflow.assert_called_once_with(
request=dict(name=WORKFLOW_NAME),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_workflows_client"))
def test_list_workflows(self, mock_client):
result = self.hook.list_workflows(
location=LOCATION,
project_id=PROJECT_ID,
filter_=FILTER_,
order_by=ORDER_BY,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.list_workflows.return_value == result
mock_client.return_value.list_workflows.assert_called_once_with(
request=dict(
parent=WORKFLOW_PARENT,
filter=FILTER_,
order_by=ORDER_BY,
),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_executions_client"))
def test_create_execution(self, mock_client):
result = self.hook.create_execution(
workflow_id=WORKFLOW_ID,
location=LOCATION,
project_id=PROJECT_ID,
execution=EXECUTION,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.create_execution.return_value == result
mock_client.return_value.create_execution.assert_called_once_with(
request=dict(
parent=EXECUTION_PARENT,
execution=EXECUTION,
),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_executions_client"))
def test_get_execution(self, mock_client):
result = self.hook.get_execution(
workflow_id=WORKFLOW_ID,
execution_id=EXECUTION_ID,
location=LOCATION,
project_id=PROJECT_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.get_execution.return_value == result
mock_client.return_value.get_execution.assert_called_once_with(
request=dict(name=EXECUTION_NAME),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_executions_client"))
def test_cancel_execution(self, mock_client):
result = self.hook.cancel_execution(
workflow_id=WORKFLOW_ID,
execution_id=EXECUTION_ID,
location=LOCATION,
project_id=PROJECT_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.cancel_execution.return_value == result
mock_client.return_value.cancel_execution.assert_called_once_with(
request=dict(name=EXECUTION_NAME),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_executions_client"))
def test_list_execution(self, mock_client):
result = self.hook.list_executions(
workflow_id=WORKFLOW_ID,
location=LOCATION,
project_id=PROJECT_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.list_executions.return_value == result
mock_client.return_value.list_executions.assert_called_once_with(
request=dict(parent=EXECUTION_PARENT),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)