def test_execute(self, mock_hook):
op = AutoMLGetModelOperator(
model_id=MODEL_ID,
location=GCP_LOCATION,
project_id=GCP_PROJECT_ID,
task_id=TASK_ID,
)
op.execute(context=None)
mock_hook.return_value.get_model.assert_called_once_with(
location=GCP_LOCATION,
metadata=None,
model_id=MODEL_ID,
project_id=GCP_PROJECT_ID,
retry=None,
timeout=None,
)
# Task dependencies created via `XComArgs`:
# create_dataset_task >> import_dataset_task
# create_dataset_task >> list_tables_spec_task
# create_dataset_task >> list_columns_spec_task
with models.DAG(
"example_gcp_get_deploy",
schedule_interval=None, # Override to match your needs
start_date=days_ago(1),
tags=["example"],
) as get_deploy_dag:
# [START howto_operator_get_model]
get_model_task = AutoMLGetModelOperator(
task_id="get_model_task",
model_id=MODEL_ID,
location=GCP_AUTOML_LOCATION,
project_id=GCP_PROJECT_ID,
)
# [END howto_operator_get_model]
# [START howto_operator_deploy_model]
deploy_model_task = AutoMLDeployModelOperator(
task_id="deploy_model_task",
model_id=MODEL_ID,
location=GCP_AUTOML_LOCATION,
project_id=GCP_PROJECT_ID,
)
# [END howto_operator_deploy_model]
with models.DAG(
"example_gcp_predict",