X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/b887df487c6cb465c3212e8ef99278350e58ff0a..5d00ecb0932f86e4d2aced3d9258b96522ef38bd:/sdk/python/tests/arvados_testutil.py diff --git a/sdk/python/tests/arvados_testutil.py b/sdk/python/tests/arvados_testutil.py index 644dfffbac..dae3dd3b7b 100644 --- a/sdk/python/tests/arvados_testutil.py +++ b/sdk/python/tests/arvados_testutil.py @@ -1,6 +1,7 @@ #!/usr/bin/env python import arvados +import contextlib import errno import hashlib import httplib @@ -8,9 +9,10 @@ import httplib2 import io import mock import os +import pycurl import Queue -import requests import shutil +import sys import tempfile import unittest @@ -43,51 +45,104 @@ def mock_responses(body, *codes, **headers): return mock.patch('httplib2.Http.request', side_effect=queue_with(( (fake_httplib2_response(code, **headers), body) for code in codes))) -# fake_requests_response, mock_get_responses and mock_put_responses -# mock calls to requests.get() and requests.put() -def fake_requests_response(code, body, **headers): - r = requests.Response() - r.status_code = code - r.reason = httplib.responses.get(code, "Unknown Response") - r.headers = headers - r.raw = io.BytesIO(body) - return r - -# The following methods patch requests.Session(), where return_value is a mock -# Session object. The put/get attributes are set on mock Session, and the -# desired put/get behavior is set on the put/get mocks. - -def mock_put_responses(body, *codes, **headers): - m = mock.MagicMock() +def mock_api_responses(api_client, body, codes, headers={}): + return mock.patch.object(api_client._http, 'request', side_effect=queue_with(( + (fake_httplib2_response(code, **headers), body) for code in codes))) + +def str_keep_locator(s): + return '{}+{}'.format(hashlib.md5(s).hexdigest(), len(s)) + +@contextlib.contextmanager +def redirected_streams(stdout=None, stderr=None): + orig_stdout, sys.stdout = sys.stdout, stdout or sys.stdout + orig_stderr, sys.stderr = sys.stderr, stderr or sys.stderr + try: + yield + finally: + sys.stdout = orig_stdout + sys.stderr = orig_stderr + + +class FakeCurl: + @classmethod + def make(cls, code, body='', headers={}): + return mock.Mock(spec=cls, wraps=cls(code, body, headers)) + + def __init__(self, code=200, body='', headers={}): + self._opt = {} + self._got_url = None + self._writer = None + self._headerfunction = None + self._resp_code = code + self._resp_body = body + self._resp_headers = headers + + def getopt(self, opt): + return self._opt.get(str(opt), None) + + def setopt(self, opt, val): + self._opt[str(opt)] = val + if opt == pycurl.WRITEFUNCTION: + self._writer = val + elif opt == pycurl.HEADERFUNCTION: + self._headerfunction = val + + def perform(self): + if not isinstance(self._resp_code, int): + raise self._resp_code + if self.getopt(pycurl.URL) is None: + raise ValueError + if self._writer is None: + raise ValueError + if self._headerfunction: + self._headerfunction("HTTP/1.1 {} Status".format(self._resp_code)) + for k, v in self._resp_headers.iteritems(): + self._headerfunction(k + ': ' + str(v)) + if type(self._resp_body) is not bool: + self._writer(self._resp_body) + + def close(self): + pass + + def reset(self): + """Prevent fake UAs from going back into the user agent pool.""" + raise Exception + + def getinfo(self, opt): + if opt == pycurl.RESPONSE_CODE: + return self._resp_code + raise Exception + +def mock_keep_responses(body, *codes, **headers): + """Patch pycurl to return fake responses and raise exceptions. + + body can be a string to return as the response body; an exception + to raise when perform() is called; or an iterable that returns a + sequence of such values. + """ + cm = mock.MagicMock() if isinstance(body, tuple): codes = list(codes) codes.insert(0, body) - m.return_value.put.side_effect = queue_with((fake_requests_response(code, b, **headers) for b, code in codes)) + responses = [ + FakeCurl.make(code=code, body=b, headers=headers) + for b, code in codes + ] else: - m.return_value.put.side_effect = queue_with((fake_requests_response(code, body, **headers) for code in codes)) - return mock.patch('requests.Session', m) - -def mock_get_responses(body, *codes, **headers): - m = mock.MagicMock() - m.return_value.get.side_effect = queue_with((fake_requests_response(code, body, **headers) for code in codes)) - return mock.patch('requests.Session', m) - -def mock_get(side_effect): - m = mock.MagicMock() - m.return_value.get.side_effect = side_effect - return mock.patch('requests.Session', m) + responses = [ + FakeCurl.make(code=code, body=body, headers=headers) + for code in codes + ] + cm.side_effect = queue_with(responses) + cm.responses = responses + return mock.patch('pycurl.Curl', cm) -def mock_put(side_effect): - m = mock.MagicMock() - m.return_value.put.side_effect = side_effect - return mock.patch('requests.Session', m) class MockStreamReader(object): def __init__(self, name='.', *data): self._name = name self._data = ''.join(data) - self._data_locators = ['{}+{}'.format(hashlib.md5(d).hexdigest(), - len(d)) for d in data] + self._data_locators = [str_keep_locator(d) for d in data] self.num_retries = 0 def name(self): @@ -104,7 +159,9 @@ class ApiClientMock(object): service_type='disk', service_host=None, service_port=None, - service_ssl_flag=False): + service_ssl_flag=False, + additional_services=[], + read_only=False): if api_mock is None: api_mock = self.api_client_mock() body = { @@ -116,7 +173,8 @@ class ApiClientMock(object): 'service_port': service_port or 65535-i, 'service_ssl_flag': service_ssl_flag, 'service_type': service_type, - } for i in range(0, count)] + 'read_only': read_only, + } for i in range(0, count)] + additional_services } self._mock_api_call(api_mock.keep_services().accessible, status, body) return api_mock