def test_assert_valid_hook_call(self, mock_hook) -> None:
task = CloudDataCatalogCreateEntryGroupOperator(
task_id="task_id",
location=TEST_LOCATION,
entry_group_id=TEST_ENTRY_GROUP_ID,
entry_group=TEST_ENTRY_GROUP,
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_group.assert_called_once_with(
location=TEST_LOCATION,
entry_group_id=TEST_ENTRY_GROUP_ID,
entry_group=TEST_ENTRY_GROUP,
project_id=TEST_PROJECT_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
ti.xcom_push.assert_called_once_with(key="entry_group_id", value=TEST_ENTRY_GROUP_ID)
self.assertEqual(result, TEST_ENTRY_GROUP_DICT)
LOCATION = "us-central1"
ENTRY_GROUP_ID = "important_data_jan_2019"
ENTRY_ID = "python_files"
TEMPLATE_ID = "template_id"
FIELD_NAME_1 = "first"
FIELD_NAME_2 = "second"
FIELD_NAME_3 = "first-rename"
with models.DAG("example_gcp_datacatalog",
start_date=days_ago(1),
schedule_interval=None) as dag:
# Create
# [START howto_operator_gcp_datacatalog_create_entry_group]
create_entry_group = CloudDataCatalogCreateEntryGroupOperator(
task_id="create_entry_group",
location=LOCATION,
entry_group_id=ENTRY_GROUP_ID,
entry_group={"display_name": "analytics data - jan 2011"},
)
# [END howto_operator_gcp_datacatalog_create_entry_group]
# [START howto_operator_gcp_datacatalog_create_entry_group_result]
create_entry_group_result = BashOperator(
task_id="create_entry_group_result",
bash_command=
"echo \"{{ task_instance.xcom_pull('create_entry_group', key='entry_group_id') }}\"",
)
# [END howto_operator_gcp_datacatalog_create_entry_group_result]
# [START howto_operator_gcp_datacatalog_create_entry_group_result2]
create_entry_group_result2 = BashOperator(
task_id="create_entry_group_result2",