def test_success_create_version(self):
with patch('airflow.gcp.operators.mlengine.MLEngineHook') \
as mock_hook:
success_response = {'name': 'some-name', 'done': True}
hook_instance = mock_hook.return_value
hook_instance.create_version.return_value = success_response
training_op = MLEngineVersionOperator(version=self.VERSION_INPUT,
**self.VERSION_DEFAULT_ARGS)
training_op.execute(None)
mock_hook.assert_called_once_with(
gcp_conn_id='google_cloud_default', delegate_to=None)
# Make sure only 'create_version' is invoked on hook instance
self.assertEqual(len(hook_instance.mock_calls), 1)
hook_instance.create_version.assert_called_once_with(
'test-project', 'test-model', self.VERSION_INPUT)
def test_success_create_version(self, mock_hook):
success_response = {'name': 'some-name', 'done': True}
hook_instance = mock_hook.return_value
hook_instance.create_version.return_value = success_response
training_op = MLEngineVersionOperator(version=TEST_VERSION,
**self.VERSION_DEFAULT_ARGS)
training_op.execute(None)
mock_hook.assert_called_once_with(gcp_conn_id='google_cloud_default',
delegate_to=None)
# Make sure only 'create_version' is invoked on hook instance
self.assertEqual(len(hook_instance.mock_calls), 1)
hook_instance.create_version.assert_called_once_with(
project_id='test-project',
model_name='test-model',
version_spec=TEST_VERSION)
model={
"name": MODEL_NAME,
})
get_model_result = BashOperator(
bash_command="echo \"{{ task_instance.xcom_pull('get-model') }}\"",
task_id="get-model-result",
)
create_version = MLEngineVersionOperator(
task_id="create-version",
project_id=PROJECT_ID,
model_name=MODEL_NAME,
operation='create',
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 = MLEngineVersionOperator(task_id="create-version-2",
project_id=PROJECT_ID,
model_name=MODEL_NAME,
operation='create',
version={
"name": "v2",
"description":
"Second version",