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,
