def test_create_version_failed_state(self, mock_mlengine_client,
mock_kfp_context, mock_dump_json,
mock_display):
version = {
'name': 'mock_version',
'description': 'the mock version',
'deploymentUri': 'gs://test-location'
}
pending_version = {'state': 'CREATING'}
pending_version.update(version)
failed_version = {
'state': 'FAILED',
'errorMessage': 'something bad happens'
}
failed_version.update(version)
mock_mlengine_client().get_version.side_effect = [
pending_version, failed_version
]
with self.assertRaises(RuntimeError) as context:
create_version('projects/mock_project/models/mock_model',
version=version,
replace_existing=True,
wait_interval=0)
self.assertEqual('Version is in failed state: something bad happens',
str(context.exception))
def test_create_version_fail(self, mock_mlengine_client, mock_kfp_context,
mock_dump_json, mock_display):
version = {
'name': 'mock_version',
'description': 'the mock version',
'deploymentUri': 'gs://test-location'
}
mock_mlengine_client().get_version.return_value = None
mock_mlengine_client().create_version.return_value = {
'name': 'mock_operation_name'
}
mock_mlengine_client().get_operation.return_value = {
'done': True,
'error': {
'code': 400,
'message': 'bad request'
}
}
with self.assertRaises(RuntimeError) as context:
create_version('projects/mock_project/models/mock_model',
version=version,
replace_existing=True,
wait_interval=30)
self.assertEqual(
'Failed to complete create version operation mock_operation_name: 400 bad request',
str(context.exception))
def test_create_version_conflict_version_fail(self, mock_mlengine_client,
mock_kfp_context,
mock_dump_json,
mock_display):
version = {
'name': 'mock_version',
'description': 'the mock version',
'deploymentUri': 'gs://test-location'
}
conflicting_version = {
'name': 'mock_version',
'description': 'the changed mock version',
'deploymentUri': 'gs://changed-test-location',
'state': 'READY'
}
mock_mlengine_client().get_version.return_value = conflicting_version
with self.assertRaises(RuntimeError) as context:
create_version('projects/mock_project/models/mock_model',
version=version,
replace_existing=False,
wait_interval=0)
self.assertEqual(
'Existing version conflicts with the name of the new version.',
str(context.exception))
def test_create_version_conflict_version_replace_succeed(
self, mock_mlengine_client, mock_kfp_context, mock_dump_json,
mock_display):
version = {
'name': 'mock_version',
'description': 'the mock version',
'deploymentUri': 'gs://test-location'
}
conflicting_version = {
'name': 'mock_version',
'description': 'the changed mock version',
'deploymentUri': 'gs://changed-test-location',
'state': 'READY'
}
mock_mlengine_client().get_version.return_value = conflicting_version
mock_mlengine_client().delete_version.return_value = {
'name': 'delete_operation_name'
}
mock_mlengine_client().create_version.return_value = {
'name': 'create_operation_name'
}
delete_operation = {'response': {}, 'done': True}
create_operation = {'response': version, 'done': True}
mock_mlengine_client().get_operation.side_effect = [
delete_operation, create_operation
]
result = create_version('projects/mock_project/models/mock_model',
version=version,
replace_existing=True,
wait_interval=0)
self.assertEqual(version, result)
def test_create_version_succeed(self, mock_mlengine_client,
mock_kfp_context, mock_dump_json,
mock_display):
version = {'description': 'the mock version'}
mock_mlengine_client().get_version.return_value = None
mock_mlengine_client().create_version.return_value = {
'name': 'mock_operation_name'
}
mock_mlengine_client().get_operation.return_value = {
'done': True,
'response': version
}
result = create_version('projects/mock_project/models/mock_model',
deployemnt_uri='gs://test-location',
version_id='mock_version',
version=version,
replace_existing=True)
self.assertEqual(version, result)
def test_create_version_dup_version_succeed(self, mock_mlengine_client,
mock_kfp_context,
mock_dump_json, mock_display):
version = {
'name': 'mock_version',
'description': 'the mock version',
'deploymentUri': 'gs://test-location'
}
pending_version = {'state': 'CREATING'}
pending_version.update(version)
ready_version = {'state': 'READY'}
ready_version.update(version)
mock_mlengine_client().get_version.side_effect = [
pending_version, ready_version
]
result = create_version('projects/mock_project/models/mock_model',
version=version,
replace_existing=True,
wait_interval=0)
self.assertEqual(ready_version, result)