def test_execute(self, mock_hook):
filter_ = "filter"
page_size = 42
op = AutoMLTablesListTableSpecsOperator(
dataset_id=DATASET_ID,
location=GCP_LOCATION,
project_id=GCP_PROJECT_ID,
filter_=filter_,
page_size=page_size,
task_id=TASK_ID,
)
op.execute(context=None)
mock_hook.return_value.list_table_specs.assert_called_once_with(
dataset_id=DATASET_ID,
filter_=filter_,
location=GCP_LOCATION,
metadata=None,
page_size=page_size,
project_id=GCP_PROJECT_ID,
retry=None,
timeout=None,
)
MODEL["dataset_id"] = dataset_id
# [START howto_operator_automl_import_data]
import_dataset_task = AutoMLImportDataOperator(
task_id="import_dataset_task",
dataset_id=dataset_id,
location=GCP_AUTOML_LOCATION,
input_config=IMPORT_INPUT_CONFIG,
)
# [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]
