def test_execute(self, mock_hook): dataset = copy.deepcopy(DATASET) dataset["name"] = DATASET_ID op = AutoMLTablesUpdateDatasetOperator( dataset=dataset, update_mask=MASK, location=GCP_LOCATION, task_id=TASK_ID, ) op.execute(context=None) mock_hook.return_value.update_dataset.assert_called_once_with( dataset=dataset, metadata=None, retry=None, timeout=None, update_mask=MASK, )
"{{ 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, location=GCP_AUTOML_LOCATION, ) # [END howto_operator_automl_update_dataset] # [START howto_operator_automl_create_model] create_model_task = AutoMLTrainModelOperator( task_id="create_model_task", model=MODEL, location=GCP_AUTOML_LOCATION, project_id=GCP_PROJECT_ID, ) model_id = create_model_task.output['model_id'] # [END howto_operator_automl_create_model]
"{{ 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, location=GCP_AUTOML_LOCATION, project_id=GCP_PROJECT_ID, ) # [END howto_operator_automl_update_dataset] # [START howto_operator_automl_create_model] create_model_task = AutoMLTrainModelOperator( task_id="create_model_task", model=MODEL, location=GCP_AUTOML_LOCATION, project_id=GCP_PROJECT_ID, ) model_id = "{{ task_instance.xcom_pull('create_model_task', key='model_id') }}" # [END howto_operator_automl_create_model]