def test_success_get_model(self, mock_hook): task = MLEngineGetModelOperator(task_id="task-id", project_id=TEST_PROJECT_ID, model_name=TEST_MODEL_NAME, gcp_conn_id=TEST_GCP_CONN_ID, delegate_to=TEST_DELEGATE_TO) result = task.execute(None) mock_hook.assert_called_once_with(delegate_to=TEST_DELEGATE_TO, gcp_conn_id=TEST_GCP_CONN_ID) mock_hook.return_value.get_model.assert_called_once_with( project_id=TEST_PROJECT_ID, model_name=TEST_MODEL_NAME) self.assertEqual(mock_hook.return_value.get_model.return_value, result)
def test_success_get_model(self, mock_hook): task = MLEngineGetModelOperator( task_id="task-id", project_id=TEST_PROJECT_ID, model_name=TEST_MODEL_NAME, gcp_conn_id=TEST_GCP_CONN_ID, delegate_to=TEST_DELEGATE_TO, impersonation_chain=TEST_IMPERSONATION_CHAIN, ) result = task.execute(None) mock_hook.assert_called_once_with( delegate_to=TEST_DELEGATE_TO, gcp_conn_id=TEST_GCP_CONN_ID, impersonation_chain=TEST_IMPERSONATION_CHAIN, ) mock_hook.return_value.get_model.assert_called_once_with( project_id=TEST_PROJECT_ID, model_name=TEST_MODEL_NAME) assert mock_hook.return_value.get_model.return_value == result
# [END howto_operator_gcp_mlengine_training] # [START howto_operator_gcp_mlengine_create_model] create_model = MLEngineCreateModelOperator( task_id="create-model", project_id=PROJECT_ID, model={ "name": MODEL_NAME, }, ) # [END howto_operator_gcp_mlengine_create_model] # [START howto_operator_gcp_mlengine_get_model] get_model = MLEngineGetModelOperator( task_id="get-model", project_id=PROJECT_ID, model_name=MODEL_NAME, ) # [END howto_operator_gcp_mlengine_get_model] # [START howto_operator_gcp_mlengine_print_model] get_model_result = BashOperator( bash_command=f"echo {get_model.output}", task_id="get-model-result", ) # [END howto_operator_gcp_mlengine_print_model] # [START howto_operator_gcp_mlengine_create_version1] create_version = MLEngineCreateVersionOperator( task_id="create-version", project_id=PROJECT_ID,