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['body'], ))
return native_request_method(*args, **kwargs)
http.request = _request_wrapper
service_mock = build('ml', 'v1', http=http)
with mock.patch.object(hook.MLEngineHook,
'get_conn',
return_value=service_mock):
return hook.MLEngineHook()
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()
from unittest import mock
from urllib.parse import urlparse, parse_qsl
import json
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.contrib.hooks import gcp_mlengine_hook 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.
