def test_execute(self, mock_hook):
op = AutoMLPredictOperator(
model_id=MODEL_ID,
location=GCP_LOCATION,
project_id=GCP_PROJECT_ID,
payload=PAYLOAD,
task_id=TASK_ID,
)
op.execute(context=None)
mock_hook.return_value.predict.assert_called_once_with(
location=GCP_LOCATION,
metadata=None,
model_id=MODEL_ID,
params={},
payload=PAYLOAD,
project_id=GCP_PROJECT_ID,
retry=None,
timeout=None,
)
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
location=GCP_AUTOML_LOCATION,
project_id=GCP_PROJECT_ID,
)
# [END howto_operator_prediction]
# [START howto_operator_batch_prediction]
batch_predict_task = AutoMLBatchPredictOperator(
task_id="batch_predict_task",
model_id=MODEL_ID,
input_config={}, # Add your config
output_config={}, # Add your config
location=GCP_AUTOML_LOCATION,
project_id=GCP_PROJECT_ID,
)
# [END howto_operator_batch_prediction]
