def test_execute(self, mock_hook):
image_detection_metadata = {}
op = AutoMLDeployModelOperator(
model_id=MODEL_ID,
image_detection_metadata=image_detection_metadata,
location=GCP_LOCATION,
project_id=GCP_PROJECT_ID,
task_id=TASK_ID,
)
op.execute(context=None)
mock_hook.return_value.deploy_model.assert_called_once_with(
image_detection_metadata={},
location=GCP_LOCATION,
metadata=None,
model_id=MODEL_ID,
project_id=GCP_PROJECT_ID,
retry=None,
timeout=None,
)
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",
schedule_interval=None, # Override to match your needs
start_date=days_ago(1),
tags=["example"],
) as predict_dag:
# [START howto_operator_prediction]
predict_task = AutoMLPredictOperator(
task_id="predict_task",
model_id=MODEL_ID,
payload={}, # Add your own payload, the used model_id must be deployed
