def test_execute(self, mock_extract_object_id, mock_create_model, mock_xcom): op = AutoMLTrainModelOperator( model=MODEL, location=GCP_LOCATION, project_id=GCP_PROJECT_ID, task_id=TASK_ID, ) op.execute(context=None) mock_create_model.assert_called_once_with( model=MODEL, location=GCP_LOCATION, project_id=GCP_PROJECT_ID, retry=None, timeout=None, metadata=None, ) mock_xcom.assert_called_once_with(None, key="model_id", value=MODEL_ID)
dataset_id = ( '{{ task_instance.xcom_pull("create_dataset_task", key="dataset_id") }}' ) import_dataset_task = AutoMLImportDataOperator( task_id="import_dataset_task", dataset_id=dataset_id, location=GCP_AUTOML_LOCATION, input_config=IMPORT_INPUT_CONFIG, ) MODEL["dataset_id"] = dataset_id create_model = AutoMLTrainModelOperator(task_id="create_model", model=MODEL, location=GCP_AUTOML_LOCATION) model_id = "{{ task_instance.xcom_pull('create_model', key='model_id') }}" delete_model_task = AutoMLDeleteModelOperator( task_id="delete_model_task", model_id=model_id, location=GCP_AUTOML_LOCATION, project_id=GCP_PROJECT_ID, ) delete_datasets_task = AutoMLDeleteDatasetOperator( task_id="delete_datasets_task", dataset_id=dataset_id, location=GCP_AUTOML_LOCATION,
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] # [START howto_operator_automl_delete_model] delete_model_task = AutoMLDeleteModelOperator( task_id="delete_model_task", model_id=model_id, location=GCP_AUTOML_LOCATION, project_id=GCP_PROJECT_ID, ) # [END howto_operator_automl_delete_model]