def test_execute(self, mock_hook):
mock_hook.return_value.batch_predict.return_value.result.return_value = BatchPredictResult(
)
mock_hook.return_value.extract_object_id = extract_object_id
op = AutoMLBatchPredictOperator(
model_id=MODEL_ID,
location=GCP_LOCATION,
project_id=GCP_PROJECT_ID,
input_config=INPUT_CONFIG,
output_config=OUTPUT_CONFIG,
task_id=TASK_ID,
prediction_params={},
)
op.execute(context=None)
mock_hook.return_value.batch_predict.assert_called_once_with(
input_config=INPUT_CONFIG,
location=GCP_LOCATION,
metadata=None,
model_id=MODEL_ID,
output_config=OUTPUT_CONFIG,
params={},
project_id=GCP_PROJECT_ID,
retry=None,
timeout=None,
)
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]