sdk/python/tests/arvados_testutil.py

   1 #!/usr/bin/env python
   2
   3 import arvados
   4 import errno
   5 import hashlib
   6 import httplib
   7 import httplib2
   8 import io
   9 import mock
  10 import os
  11 import requests
  12 import shutil
  13 import tempfile
  14 import unittest
  15
  16 # Use this hostname when you want to make sure the traffic will be
  17 # instantly refused.  100::/64 is a dedicated black hole.
  18 TEST_HOST = '100::'
  19
  20 skip_sleep = mock.patch('time.sleep', lambda n: None)  # clown'll eat me
  21
  22 # fake_httplib2_response and mock_responses
  23 # mock calls to httplib2.Http.request()
  24 def fake_httplib2_response(code, **headers):
  25     headers.update(status=str(code),
  26                    reason=httplib.responses.get(code, "Unknown Response"))
  27     return httplib2.Response(headers)
  28
  29 def mock_responses(body, *codes, **headers):
  30     return mock.patch('httplib2.Http.request', side_effect=(
  31             (fake_httplib2_response(code, **headers), body) for code in codes))
  32
  33 # fake_requests_response, mock_get_responses and mock_put_responses
  34 # mock calls to requests.get() and requests.put()
  35 def fake_requests_response(code, body, **headers):
  36     r = requests.Response()
  37     r.status_code = code
  38     r.reason = httplib.responses.get(code, "Unknown Response")
  39     r.headers = headers
  40     r.raw = io.BytesIO(body)
  41     return r
  42
  43 def mock_get_responses(body, *codes, **headers):
  44     return mock.patch('requests.get', side_effect=(
  45         fake_requests_response(code, body, **headers) for code in codes))
  46
  47 def mock_put_responses(body, *codes, **headers):
  48     return mock.patch('requests.put', side_effect=(
  49         fake_requests_response(code, body, **headers) for code in codes))
  50
  51 def mock_requestslib_responses(method, body, *codes, **headers):
  52     return mock.patch(method, side_effect=(
  53         fake_requests_response(code, body, **headers) for code in codes))
  54
  55 class MockStreamReader(object):
  56     def __init__(self, name='.', *data):
  57         self._name = name
  58         self._data = ''.join(data)
  59         self._data_locators = ['{}+{}'.format(hashlib.md5(d).hexdigest(),
  60                                               len(d)) for d in data]
  61         self.num_retries = 0
  62
  63     def name(self):
  64         return self._name
  65
  66     def readfrom(self, start, size, num_retries=None):
  67         return self._data[start:start + size]
  68
  69
  70 class ApiClientMock(object):
  71     def api_client_mock(self):
  72         return mock.MagicMock(name='api_client_mock')
  73
  74     def mock_keep_services(self, api_mock=None, status=200, count=12,
  75                            service_type='disk',
  76                            service_host=None,
  77                            service_port=None,
  78                            service_ssl_flag=False):
  79         if api_mock is None:
  80             api_mock = self.api_client_mock()
  81         body = {
  82             'items_available': count,
  83             'items': [{
  84                 'uuid': 'zzzzz-bi6l4-{:015x}'.format(i),
  85                 'owner_uuid': 'zzzzz-tpzed-000000000000000',
  86                 'service_host': service_host or 'keep0x{:x}'.format(i),
  87                 'service_port': service_port or 65535-i,
  88                 'service_ssl_flag': service_ssl_flag,
  89                 'service_type': service_type,
  90             } for i in range(0, count)]
  91         }
  92         self._mock_api_call(api_mock.keep_services().accessible, status, body)
  93         return api_mock
  94
  95     def _mock_api_call(self, mock_method, code, body):
  96         mock_method = mock_method().execute
  97         if code == 200:
  98             mock_method.return_value = body
  99         else:
 100             mock_method.side_effect = arvados.errors.ApiError(
 101                 fake_httplib2_response(code), "{}")
 102
 103
 104 class ArvadosBaseTestCase(unittest.TestCase):
 105     # This class provides common utility functions for our tests.
 106
 107     def setUp(self):
 108         self._tempdirs = []
 109
 110     def tearDown(self):
 111         for workdir in self._tempdirs:
 112             shutil.rmtree(workdir, ignore_errors=True)
 113
 114     def make_tmpdir(self):
 115         self._tempdirs.append(tempfile.mkdtemp())
 116         return self._tempdirs[-1]
 117
 118     def data_file(self, filename):
 119         try:
 120             basedir = os.path.dirname(__file__)
 121         except NameError:
 122             basedir = '.'
 123         return open(os.path.join(basedir, 'data', filename))
 124
 125     def build_directory_tree(self, tree):
 126         tree_root = self.make_tmpdir()
 127         for leaf in tree:
 128             path = os.path.join(tree_root, leaf)
 129             try:
 130                 os.makedirs(os.path.dirname(path))
 131             except OSError as error:
 132                 if error.errno != errno.EEXIST:
 133                     raise
 134             with open(path, 'w') as tmpfile:
 135                 tmpfile.write(leaf)
 136         return tree_root
 137
 138     def make_test_file(self, text="test"):
 139         testfile = tempfile.NamedTemporaryFile()
 140         testfile.write(text)
 141         testfile.flush()
 142         return testfile