def test_delete_model_remote_runner_cancel(self, mock_execution_context, mock_post_requests, mock_delete_requests, _, mock_auth): creds = mock.Mock() creds.token = 'fake_token' mock_auth.return_value = [creds, 'project'] delete_model_lro = mock.Mock() delete_model_lro.json.return_value = { 'name': self._lro_name, 'done': True, } mock_delete_requests.return_value = delete_model_lro mock_execution_context.return_value = None delete_model_remote_runner.delete_model(self._type, '', '', self._payload, self._gcp_resources_path) # Call cancellation handler mock_execution_context.call_args[1]['on_cancel']() mock_post_requests.assert_called_once_with( url=f'{self._uri_prefix}{self._lro_name}:cancel', data='', headers={ 'Content-type': 'application/json', 'Authorization': 'Bearer fake_token', })
def test_delete_model_remote_runner_succeeded(self, mock_delete_requests, _, mock_auth): creds = mock.Mock() creds.token = 'fake_token' mock_auth.return_value = [creds, 'project'] delete_model_lro = mock.Mock() delete_model_lro.json.return_value = { 'name': self._lro_name, 'done': True, } mock_delete_requests.return_value = delete_model_lro delete_model_remote_runner.delete_model(self._type, '', '', self._payload, self._gcp_resources_path) mock_delete_requests.assert_called_once_with( url= f'{self._uri_prefix}projects/test_project/locations/test_region/models/m12', data='', 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_delete_model_remote_runner_poll_till_succeeded( self, mock_time_sleep, mock_get_requests, mock_delete_requests, _, mock_auth): creds = mock.Mock() creds.token = 'fake_token' mock_auth.return_value = [creds, 'project'] delete_model_lro = mock.Mock() delete_model_lro.json.return_value = { 'name': self._lro_name, 'done': False } mock_delete_requests.return_value = delete_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 delete_model_remote_runner.delete_model(self._type, '', '', self._payload, self._gcp_resources_path) self.assertEqual(mock_delete_requests.call_count, 1) self.assertEqual(mock_time_sleep.call_count, 2) self.assertEqual(mock_get_requests.call_count, 2)
def test_delete_model_remote_runner_raises_exception_on_internal_error( self, mock_delete_requests, _, mock_auth): creds = mock.Mock() creds.token = 'fake_token' mock_auth.return_value = [creds, 'project'] delete_model_lro = mock.Mock() delete_model_lro.json.return_value = { 'name': self._lro_name, 'done': True, 'error': { 'code': 1 } } mock_delete_requests.return_value = delete_model_lro with self.assertRaises(SystemExit): delete_model_remote_runner.delete_model(self._type, '', '', self._payload, self._gcp_resources_path)