def test_assert_valid_hook_call(self, mock_hook) -> None:
task = CloudDataCatalogCreateEntryOperator(
task_id="task_id",
location=TEST_LOCATION,
entry_group=TEST_ENTRY_GROUP_ID,
entry_id=TEST_ENTRY_ID,
entry=TEST_ENTRY,
project_id=TEST_PROJECT_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
gcp_conn_id=TEST_GCP_CONN_ID,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
ti = mock.MagicMock()
result = task.execute(context={"task_instance": ti})
mock_hook.assert_called_once_with(
gcp_conn_id=TEST_GCP_CONN_ID,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
mock_hook.return_value.create_entry.assert_called_once_with(
location=TEST_LOCATION,
entry_group=TEST_ENTRY_GROUP_ID,
entry_id=TEST_ENTRY_ID,
entry=TEST_ENTRY,
project_id=TEST_PROJECT_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
ti.xcom_push.assert_called_once_with(key="entry_id", value=TEST_ENTRY_ID)
self.assertEqual(TEST_ENTRY_DICT, result)
# [START howto_operator_gcp_datacatalog_create_entry_group_result2]
create_entry_group_result2 = BashOperator(
task_id="create_entry_group_result2",
bash_command=
"echo \"{{ task_instance.xcom_pull('create_entry_group') }}\"",
)
# [END howto_operator_gcp_datacatalog_create_entry_group_result2]
# [START howto_operator_gcp_datacatalog_create_entry_gcs]
create_entry_gcs = CloudDataCatalogCreateEntryOperator(
task_id="create_entry_gcs",
location=LOCATION,
entry_group=ENTRY_GROUP_ID,
entry_id=ENTRY_ID,
entry={
"display_name": "Wizard",
"type": "FILESET",
"gcs_fileset_spec": {
"file_patterns": ["gs://test-datacatalog/**"]
},
},
)
# [END howto_operator_gcp_datacatalog_create_entry_gcs]
# [START howto_operator_gcp_datacatalog_create_entry_gcs_result]
create_entry_gcs_result = BashOperator(
task_id="create_entry_gcs_result",
bash_command=
"echo \"{{ task_instance.xcom_pull('create_entry_gcs', key='entry_id') }}\"",
)
# [END howto_operator_gcp_datacatalog_create_entry_gcs_result]