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