def test_success(self, mock_hook):
task = MLEngineCreateVersionOperator(
task_id="task-id",
project_id=TEST_PROJECT_ID,
model_name=TEST_MODEL_NAME,
version=TEST_VERSION,
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.create_version.assert_called_once_with(
project_id=TEST_PROJECT_ID,
model_name=TEST_MODEL_NAME,
version_spec=TEST_VERSION)
def test_missing_version(self):
with self.assertRaises(AirflowException):
MLEngineCreateVersionOperator(
task_id="task-id",
project_id=TEST_PROJECT_ID,
model_name=TEST_MODEL_NAME,
version=None,
gcp_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
)
model={
"name": MODEL_NAME,
})
get_model_result = BashOperator(
bash_command="echo \"{{ task_instance.xcom_pull('get-model') }}\"",
task_id="get-model-result",
)
create_version = MLEngineCreateVersionOperator(
task_id="create-version",
project_id=PROJECT_ID,
model_name=MODEL_NAME,
version={
"name": "v1",
"description": "First-version",
"deployment_uri": '{}/keras_export/'.format(JOB_DIR),
"runtime_version": "1.14",
"machineType": "mls1-c1-m2",
"framework": "TENSORFLOW",
"pythonVersion": "3.5"
})
create_version_2 = MLEngineCreateVersionOperator(
task_id="create-version-2",
project_id=PROJECT_ID,
model_name=MODEL_NAME,
version={
"name": "v2",
"description": "Second version",
"deployment_uri": SAVED_MODEL_PATH,