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.