def test_success(self, mock_hook):
task = MLEngineDeleteVersionOperator(
task_id="task-id",
project_id=TEST_PROJECT_ID,
model_name=TEST_MODEL_NAME,
version_name=TEST_VERSION_NAME,
gcp_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
)
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.delete_version.assert_called_once_with(
project_id=TEST_PROJECT_ID,
model_name=TEST_MODEL_NAME,
version_name=TEST_VERSION_NAME
)
def test_missing_model_name(self):
with pytest.raises(AirflowException):
MLEngineDeleteVersionOperator(
task_id="task-id",
project_id=TEST_PROJECT_ID,
model_name=None,
version_name=TEST_VERSION_NAME,
gcp_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
)
prediction = MLEngineStartBatchPredictionJobOperator(
task_id="prediction",
project_id=PROJECT_ID,
job_id="prediction-{{ ts_nodash }}-{{ params.model_name }}",
region="us-central1",
model_name=MODEL_NAME,
data_format="TEXT",
input_paths=[PREDICTION_INPUT],
output_path=PREDICTION_OUTPUT,
labels={"job_type": "prediction"},
)
# [END howto_operator_gcp_mlengine_get_prediction]
# [START howto_operator_gcp_mlengine_delete_version]
delete_version = MLEngineDeleteVersionOperator(task_id="delete-version",
project_id=PROJECT_ID,
model_name=MODEL_NAME,
version_name="v1")
# [END howto_operator_gcp_mlengine_delete_version]
# [START howto_operator_gcp_mlengine_delete_model]
delete_model = MLEngineDeleteModelOperator(task_id="delete-model",
project_id=PROJECT_ID,
model_name=MODEL_NAME,
delete_contents=True)
# [END howto_operator_gcp_mlengine_delete_model]
training >> create_version
training >> create_version_2
create_model >> get_model >> [get_model_result, delete_model]
create_model >> get_model >> delete_model
create_model >> create_version >> create_version_2 >> set_defaults_version >> list_version