def test_execute(self, mock_hook):
table_spec = "table_spec_id"
filter_ = "filter"
page_size = 42
op = AutoMLTablesListColumnSpecsOperator(
dataset_id=DATASET_ID,
table_spec_id=table_spec,
location=GCP_LOCATION,
project_id=GCP_PROJECT_ID,
field_mask=MASK,
filter_=filter_,
page_size=page_size,
task_id=TASK_ID,
)
op.execute(context=None)
mock_hook.return_value.list_column_specs.assert_called_once_with(
dataset_id=DATASET_ID,
field_mask=MASK,
filter_=filter_,
location=GCP_LOCATION,
metadata=None,
page_size=page_size,
project_id=GCP_PROJECT_ID,
retry=None,
table_spec_id=table_spec,
timeout=None,
)
# [END howto_operator_automl_import_data]
# [START howto_operator_automl_specs]
list_tables_spec_task = AutoMLTablesListTableSpecsOperator(
task_id="list_tables_spec_task",
dataset_id=dataset_id,
location=GCP_AUTOML_LOCATION,
project_id=GCP_PROJECT_ID,
)
# [END howto_operator_automl_specs]
# [START howto_operator_automl_column_specs]
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,
)
# [END howto_operator_automl_column_specs]
# [START howto_operator_automl_update_dataset]
update = deepcopy(DATASET)
update[
"name"] = '{{ task_instance.xcom_pull("create_dataset_task")["name"] }}'
update["tables_dataset_metadata"][ # type: ignore
"target_column_spec_id"] = "{{ get_target_column_spec(task_instance.xcom_pull('list_columns_spec_task'), target) }}"
update_dataset_task = AutoMLTablesUpdateDatasetOperator(
task_id="update_dataset_task",
dataset=update,