def test_execute(self, mock_hook, mock_xcom):
op = AutoMLListDatasetOperator(location=GCP_LOCATION, project_id=GCP_PROJECT_ID, task_id=TASK_ID)
op.execute(context=None)
mock_hook.return_value.list_datasets.assert_called_once_with(
location=GCP_LOCATION,
metadata=None,
project_id=GCP_PROJECT_ID,
retry=None,
timeout=None,
)
mock_xcom.assert_called_once_with(None, key="dataset_id_list", value=[])
project_id=GCP_PROJECT_ID,
)
list_columns_spec_task = AutoMLTablesListColumnSpecsOperator(
task_id="list_columns_spec_task",
dataset_id=dataset_id,
table_spec_id=
"{{ extract_object_id(task_instance.xcom_pull('list_tables_spec_task')[0]) }}",
location=GCP_AUTOML_LOCATION,
project_id=GCP_PROJECT_ID,
)
# [START howto_operator_list_dataset]
list_datasets_task = AutoMLListDatasetOperator(
task_id="list_datasets_task",
location=GCP_AUTOML_LOCATION,
project_id=GCP_PROJECT_ID,
)
# [END howto_operator_list_dataset]
# [START howto_operator_delete_dataset]
delete_datasets_task = AutoMLDeleteDatasetOperator(
task_id="delete_datasets_task",
dataset_id=
"{{ task_instance.xcom_pull('list_datasets_task', key='dataset_id_list') | list }}",
location=GCP_AUTOML_LOCATION,
project_id=GCP_PROJECT_ID,
)
# [END howto_operator_delete_dataset]
(import_dataset_task >> list_tables_spec_task >> list_columns_spec_task >>
