def setUp(self) -> None: super().setUp() with mock.patch( 'airflow.gcp.hooks.mlengine.MLEngineHook.__init__', new=mock_base_gcp_hook_no_default_project_id, ): self.hook = hook.MLEngineHook()
def test_create_model(self, mock_get_conn): project_id = 'test-project' model_name = 'test-model' model = { 'name': model_name, } project_path = 'projects/{}'.format(project_id) ( mock_get_conn.return_value. projects.return_value. models.return_value. create.return_value. execute.return_value ) = model mle_engine_hook = hook.MLEngineHook() create_model_response = mle_engine_hook.create_model( project_id=project_id, model=model ) self.assertEqual(create_model_response, model) mock_get_conn.assert_has_calls([ mock.call().projects().models().create(body=model, parent=project_path), mock.call().projects().models().create().execute() ])
def test_set_default_version(self, mock_get_conn): project_id = 'test-project' model_name = 'test-model' version_name = 'test-version' operation_path = 'projects/{}/operations/test-operation'.format(project_id) version_path = 'projects/{}/models/{}/versions/{}'.format(project_id, model_name, version_name) operation_done = {'name': operation_path, 'done': True} ( mock_get_conn.return_value. projects.return_value. models.return_value. versions.return_value. setDefault.return_value. execute.return_value ) = operation_done mle_engine_hook = hook.MLEngineHook() set_default_version_response = mle_engine_hook.set_default_version( project_id=project_id, model_name=model_name, version_name=version_name ) self.assertEqual(set_default_version_response, operation_done) mock_get_conn.assert_has_calls([ mock.call().projects().models().versions().setDefault(body={}, name=version_path), mock.call().projects().models().versions().setDefault().execute() ], any_order=True)
def test_delete_model_when_not_exists(self, mock_get_conn, mock_log): project_id = 'test-project' model_name = 'test-model' model_path = 'projects/{}/models/{}'.format(project_id, model_name) http_error = HttpError( resp=mock.MagicMock(status=404, reason="Model not found."), content=b'Model not found.' ) ( mock_get_conn.return_value. projects.return_value. models.return_value. delete.return_value. execute.side_effect ) = [http_error] mle_engine_hook = hook.MLEngineHook() mle_engine_hook.delete_model( project_id=project_id, model_name=model_name ) mock_get_conn.assert_has_calls([ mock.call().projects().models().delete(name=model_path), mock.call().projects().models().delete().execute() ]) mock_log.error.assert_called_once_with('Model was not found: %s', http_error)
def test_mle_engine_client_creation(self, mock_build, mock_authorize): mle_engine_hook = hook.MLEngineHook() result = mle_engine_hook.get_conn() mock_build.assert_called_with('ml', 'v1', http=mock_authorize.return_value, cache_discovery=False) self.assertEqual(mock_build.return_value, result)
def test_delete_model_with_contents(self, mock_get_conn): project_id = 'test-project' model_name = 'test-model' model_path = 'projects/{}/models/{}'.format(project_id, model_name) operation_path = 'projects/{}/operations/test-operation'.format(project_id) operation_done = {'name': operation_path, 'done': True} version_names = ["AAA", "BBB", "CCC"] versions = [{ 'name': 'projects/{}/models/{}/versions/{}'.format(project_id, model_name, version_name), "isDefault": i == 0 } for i, version_name in enumerate(version_names)] ( mock_get_conn.return_value. projects.return_value. operations.return_value. get.return_value. execute.return_value ) = operation_done ( mock_get_conn.return_value. projects.return_value. models.return_value. versions.return_value. list.return_value. execute.return_value ) = {"versions": versions} ( mock_get_conn.return_value. projects.return_value. models.return_value. versions.return_value. list_next.return_value ) = None mle_engine_hook = hook.MLEngineHook() mle_engine_hook.delete_model( project_id=project_id, model_name=model_name, delete_contents=True ) mock_get_conn.assert_has_calls( [ mock.call().projects().models().delete(name=model_path), mock.call().projects().models().delete().execute() ] + [ mock.call().projects().models().versions().delete( name='projects/{}/models/{}/versions/{}'.format(project_id, model_name, version_name), ) for version_name in version_names ], any_order=True )
def __enter__(self): http = HttpMockSequence(self._responses) native_request_method = http.request # Collecting requests to validate at __exit__. def _request_wrapper(*args, **kwargs): self._actual_requests.append(args + (kwargs.get('body', ''), )) return native_request_method(*args, **kwargs) http.request = _request_wrapper discovery = requests.get( 'https://www.googleapis.com/discovery/v1/apis/ml/v1/rest') service_mock = build_from_document(discovery.json(), http=http) with mock.patch.object(hook.MLEngineHook, 'get_conn', return_value=service_mock): return hook.MLEngineHook()
def test_create_mlengine_job_check_existing_job_failed(self, mock_get_conn): project_id = 'test-project' job_id = 'test-job-id' my_job = { 'jobId': job_id, 'foo': 4815162342, 'state': 'SUCCEEDED', 'someInput': { 'input': 'someInput' } } different_job = { 'jobId': job_id, 'foo': 4815162342, 'state': 'SUCCEEDED', 'someInput': { 'input': 'someDifferentInput' } } error_job_exists = HttpError(resp=mock.MagicMock(status=409), content=b'Job already exists') ( mock_get_conn.return_value. projects.return_value. jobs.return_value. create.return_value. execute.side_effect ) = error_job_exists ( mock_get_conn.return_value. projects.return_value. jobs.return_value. get.return_value. execute.return_value ) = different_job def check_input(existing_job): return existing_job.get('someInput', None) == \ my_job['someInput'] with self.assertRaises(HttpError): mle_engine_hook = hook.MLEngineHook() mle_engine_hook.create_job( project_id=project_id, job=my_job, use_existing_job_fn=check_input)
def test_create_mlengine_job(self, mock_get_conn): project_id = 'test-project' job_id = 'test-job-id' project_path = 'projects/{}'.format(project_id) job_path = 'projects/{}/jobs/{}'.format(project_id, job_id) new_job = { 'jobId': job_id, 'foo': 4815162342, } job_succeeded = { 'jobId': job_id, 'state': 'SUCCEEDED', } job_queued = { 'jobId': job_id, 'state': 'QUEUED', } ( mock_get_conn.return_value. projects.return_value. jobs.return_value. create.return_value. execute.return_value ) = job_queued ( mock_get_conn.return_value. projects.return_value. jobs.return_value. get.return_value. execute.side_effect ) = [job_queued, job_succeeded] mle_engine_hook = hook.MLEngineHook() create_job_response = mle_engine_hook.create_job( project_id=project_id, job=new_job ) self.assertEqual(create_job_response, job_succeeded) mock_get_conn.assert_has_calls([ mock.call().projects().jobs().create(body=new_job, parent=project_path), mock.call().projects().jobs().get(name=job_path), mock.call().projects().jobs().get().execute() ], any_order=True)
def test_delete_version(self, mock_get_conn): project_id = 'test-project' model_name = 'test-model' version_name = 'test-version' operation_path = 'projects/{}/operations/test-operation'.format(project_id) version_path = 'projects/{}/models/{}/versions/{}'.format(project_id, model_name, version_name) version = {'name': operation_path} operation_not_done = {'name': operation_path, 'done': False} operation_done = {'name': operation_path, 'done': True} ( mock_get_conn.return_value. projects.return_value. operations.return_value. get.return_value. execute.side_effect ) = [operation_not_done, operation_done] ( mock_get_conn.return_value. projects.return_value. models.return_value. versions.return_value. delete.return_value. execute.return_value ) = version mle_engine_hook = hook.MLEngineHook() delete_version_response = mle_engine_hook.delete_version( project_id=project_id, model_name=model_name, version_name=version_name) self.assertEqual(delete_version_response, operation_done) mock_get_conn.assert_has_calls([ mock.call().projects().models().versions().delete(name=version_path), mock.call().projects().models().versions().delete().execute(), mock.call().projects().operations().get(name=operation_path), mock.call().projects().operations().get().execute() ], any_order=True)
def test_delete_model(self, mock_get_conn): project_id = 'test-project' model_name = 'test-model' model = {'model': model_name} model_path = 'projects/{}/models/{}'.format(project_id, model_name) ( mock_get_conn.return_value. projects.return_value. models.return_value. delete.return_value. execute.return_value ) = model mle_engine_hook = hook.MLEngineHook() mle_engine_hook.delete_model( project_id=project_id, model_name=model_name ) mock_get_conn.assert_has_calls([ mock.call().projects().models().delete(name=model_path), mock.call().projects().models().delete().execute() ])
def test_list_versions(self, mock_get_conn): project_id = 'test-project' model_name = 'test-model' model_path = 'projects/{}/models/{}'.format(project_id, model_name) version_names = ['ver_{}'.format(ix) for ix in range(3)] response_bodies = [ { 'nextPageToken': "TOKEN-{}".format(ix), 'versions': [ver] } for ix, ver in enumerate(version_names)] response_bodies[-1].pop('nextPageToken') pages_requests = [ mock.Mock(**{'execute.return_value': body}) for body in response_bodies ] versions_mock = mock.Mock( **{'list.return_value': pages_requests[0], 'list_next.side_effect': pages_requests[1:] + [None]} ) ( mock_get_conn.return_value. projects.return_value. models.return_value. versions.return_value ) = versions_mock mle_engine_hook = hook.MLEngineHook() list_versions_response = mle_engine_hook.list_versions( project_id=project_id, model_name=model_name) self.assertEqual(list_versions_response, version_names) mock_get_conn.assert_has_calls([ mock.call().projects().models().versions().list(pageSize=100, parent=model_path), mock.call().projects().models().versions().list().execute(), ] + [ mock.call().projects().models().versions().list_next( previous_request=pages_requests[i], previous_response=response_bodies[i] ) for i in range(3) ], any_order=True)
def test_create_version(self, mock_get_conn): project_id = 'test-project' model_name = 'test-model' version_name = 'test-version' version = {'name': version_name} operation_path = 'projects/{}/operations/test-operation'.format(project_id) model_path = 'projects/{}/models/{}'.format(project_id, model_name) operation_done = {'name': operation_path, 'done': True} ( mock_get_conn.return_value. projects.return_value. models.return_value. versions.return_value. create.return_value. execute.return_value ) = version ( mock_get_conn.return_value. projects.return_value. operations.return_value. get.return_value. execute.return_value ) = {'name': operation_path, 'done': True} mle_engine_hook = hook.MLEngineHook() create_version_response = mle_engine_hook.create_version( project_id=project_id, model_name=model_name, version_spec=version ) self.assertEqual(create_version_response, operation_done) mock_get_conn.assert_has_calls([ mock.call().projects().models().versions().create(body=version, parent=model_path), mock.call().projects().models().versions().create().execute(), mock.call().projects().operations().get(name=version_name), ], any_order=True)
def test_create_mlengine_job_reuse_existing_job_by_default(self, mock_get_conn): project_id = 'test-project' job_id = 'test-job-id' project_path = 'projects/{}'.format(project_id) job_path = 'projects/{}/jobs/{}'.format(project_id, job_id) job_succeeded = { 'jobId': job_id, 'foo': 4815162342, 'state': 'SUCCEEDED', } error_job_exists = HttpError(resp=mock.MagicMock(status=409), content=b'Job already exists') ( mock_get_conn.return_value. projects.return_value. jobs.return_value. create.return_value. execute.side_effect ) = error_job_exists ( mock_get_conn.return_value. projects.return_value. jobs.return_value. get.return_value. execute.return_value ) = job_succeeded mle_engine_hook = hook.MLEngineHook() create_job_response = mle_engine_hook.create_job( project_id=project_id, job=job_succeeded) self.assertEqual(create_job_response, job_succeeded) mock_get_conn.assert_has_calls([ mock.call().projects().jobs().create(body=job_succeeded, parent=project_path), mock.call().projects().jobs().create().execute(), mock.call().projects().jobs().get(name=job_path), mock.call().projects().jobs().get().execute() ], any_order=True)
def test_get_model(self, mock_get_conn): project_id = 'test-project' model_name = 'test-model' model = {'model': model_name} model_path = 'projects/{}/models/{}'.format(project_id, model_name) ( mock_get_conn.return_value. projects.return_value. models.return_value. get.return_value. execute.return_value ) = model mle_engine_hook = hook.MLEngineHook() get_model_response = mle_engine_hook.get_model( project_id=project_id, model_name=model_name ) self.assertEqual(get_model_response, model) mock_get_conn.assert_has_calls([ mock.call().projects().models().get(name=model_path), mock.call().projects().models().get().execute() ])
def setUp(self) -> None: super().setUp() self.hook = hook.MLEngineHook()
import json import unittest from unittest import mock from urllib.parse import parse_qsl, urlparse import requests from google.auth.exceptions import GoogleAuthError from googleapiclient.discovery import build_from_document from googleapiclient.errors import HttpError from googleapiclient.http import HttpMockSequence from airflow.gcp.hooks import mlengine as hook cml_available = True try: hook.MLEngineHook().get_conn() except GoogleAuthError: cml_available = False class _TestMLEngineHook: def __init__(self, test_cls, responses, expected_requests): """ Init method. Usage example: with _TestMLEngineHook(self, responses, expected_requests) as hook: self.run_my_test(hook) Args: test_cls: The caller's instance used for test communication.