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.CloudMLHook,
'get_conn',
return_value=service_mock):
return hook.CloudMLHook()
import json
import mock
import unittest
try: # python 2
from urlparse import urlparse, parse_qsl
except ImportError: #python 3
from urllib.parse import urlparse, parse_qsl
from airflow.contrib.hooks import gcp_cloudml_hook as hook
from apiclient.discovery import build
from apiclient.http import HttpMockSequence
from oauth2client.contrib.gce import HttpAccessTokenRefreshError
cml_available = True
try:
hook.CloudMLHook().get_conn()
except HttpAccessTokenRefreshError:
cml_available = False
class _TestCloudMLHook(object):
def __init__(self, test_cls, responses, expected_requests):
"""
Init method.
Usage example:
with _TestCloudMLHook(self, responses, expected_requests) as hook:
self.run_my_test(hook)
Args:
test_cls: The caller's instance used for test communication.