def test_model_deploy_remote_runner_cancel(self, mock_execution_context,
mock_post_requests, _,
mock_auth):
creds = mock.Mock()
creds.token = 'fake_token'
mock_auth.return_value = [creds, 'project']
deploy_model_lro = mock.Mock()
deploy_model_lro.json.return_value = {
'name': self._lro_name,
'done': True,
}
mock_post_requests.return_value = deploy_model_lro
mock_execution_context.return_value = None
deploy_model_remote_runner.deploy_model(self._type, '', '',
self._payload,
self._gcp_resources_path)
# Call cancellation handler
mock_execution_context.call_args[1]['on_cancel']()
self.assertEqual(mock_post_requests.call_count, 2)
mock_post_requests.assert_called_with(
url=f'{self._uri_prefix}{self._lro_name}:cancel',
data='',
headers={
'Content-type': 'application/json',
'Authorization': 'Bearer fake_token',
})
def test_model_deploy_remote_runner_succeeded(self, mock_post_requests, _,
mock_auth):
creds = mock.Mock()
creds.token = 'fake_token'
mock_auth.return_value = [creds, 'project']
deploy_model_lro = mock.Mock()
deploy_model_lro.json.return_value = {
'name': self._lro_name,
'done': True,
}
mock_post_requests.return_value = deploy_model_lro
deploy_model_remote_runner.deploy_model(self._type, '', '',
self._payload,
self._gcp_resources_path)
mock_post_requests.assert_called_once_with(
url=
f'{self._uri_prefix}projects/test_project/locations/test_region/endpoints/e12:deployModel',
data=self._payload,
headers={
'Content-type': 'application/json',
'Authorization': 'Bearer fake_token',
'User-Agent': 'google-cloud-pipeline-components'
})
with open(self._gcp_resources_path) as f:
serialized_gcp_resources = f.read()
# Instantiate GCPResources Proto
lro_resources = json_format.Parse(serialized_gcp_resources,
GcpResources())
self.assertEqual(len(lro_resources.resources), 1)
self.assertEqual(lro_resources.resources[0].resource_uri,
self._uri_prefix + self._lro_name)
def test_deploy_model_remote_runner_poll_till_succeeded(
self, mock_time_sleep, mock_get_requests, mock_post_requests, _,
mock_auth):
creds = mock.Mock()
creds.token = 'fake_token'
mock_auth.return_value = [creds, 'project']
deploy_model_lro = mock.Mock()
deploy_model_lro.json.return_value = {
'name': self._lro_name,
'done': False
}
mock_post_requests.return_value = deploy_model_lro
poll_lro = mock.Mock()
poll_lro.json.side_effect = [{
'name': self._lro_name,
'done': False
}, {
'name': self._lro_name,
'done': True
}]
mock_get_requests.return_value = poll_lro
deploy_model_remote_runner.deploy_model(self._type, '', '',
self._payload,
self._gcp_resources_path)
self.assertEqual(mock_post_requests.call_count, 1)
self.assertEqual(mock_time_sleep.call_count, 2)
self.assertEqual(mock_get_requests.call_count, 2)
def test_deploy_model_remote_runner_raises_exception_on_internal_error(
self, mock_post_requests, _, mock_auth):
creds = mock.Mock()
creds.token = 'fake_token'
mock_auth.return_value = [creds, 'project']
deploy_model_lro = mock.Mock()
deploy_model_lro.json.return_value = {
'name': self._lro_name,
'done': True,
'error': {
'code': 1
}
}
mock_post_requests.return_value = deploy_model_lro
with self.assertRaises(SystemExit):
deploy_model_remote_runner.deploy_model(self._type, '', '',
self._payload,
self._gcp_resources_path)