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_put_responses(body, *codes, **headers):
  44     m = mock.MagicMock()
  45     if isinstance(body, tuple):
  46         codes = list(codes)
  47         codes.insert(0, body)
  48         m.return_value.put.side_effect = (fake_requests_response(code, b, **headers) for b, code in codes)
  49     else:
  50         m.return_value.put.side_effect = (fake_requests_response(code, body, **headers) for code in codes)
  51     return mock.patch('requests.Session', m)
  52
  53 def mock_get_responses(body, *codes, **headers):
  54     m = mock.MagicMock()
  55     m.return_value.get.side_effect = (fake_requests_response(code, body, **headers) for code in codes)
  56     return mock.patch('requests.Session', m)
  57
  58 def mock_get(side_effect):
  59     m = mock.MagicMock()
  60     m.return_value.get.side_effect = side_effect
  61     return mock.patch('requests.Session', m)
  62
  63 def mock_put(side_effect):
  64     m = mock.MagicMock()
  65     m.return_value.put.side_effect = side_effect
  66     return mock.patch('requests.Session', m)
  67
  68 class MockStreamReader(object):
  69     def __init__(self, name='.', *data):
  70         self._name = name
  71         self._data = ''.join(data)
  72         self._data_locators = ['{}+{}'.format(hashlib.md5(d).hexdigest(),
  73                                               len(d)) for d in data]
  74         self.num_retries = 0
  75
  76     def name(self):
  77         return self._name
  78
  79     def readfrom(self, start, size, num_retries=None):
  80         self._readfrom(start, size, num_retries=num_retries)
  81
  82     def _readfrom(self, start, size, num_retries=None):
  83         return self._data[start:start + size]
  84
  85 class ApiClientMock(object):
  86     def api_client_mock(self):
  87         return mock.MagicMock(name='api_client_mock')
  88
  89     def mock_keep_services(self, api_mock=None, status=200, count=12,
  90                            service_type='disk',
  91                            service_host=None,
  92                            service_port=None,
  93                            service_ssl_flag=False):
  94         if api_mock is None:
  95             api_mock = self.api_client_mock()
  96         body = {
  97             'items_available': count,
  98             'items': [{
  99                 'uuid': 'zzzzz-bi6l4-{:015x}'.format(i),
 100                 'owner_uuid': 'zzzzz-tpzed-000000000000000',
 101                 'service_host': service_host or 'keep0x{:x}'.format(i),
 102                 'service_port': service_port or 65535-i,
 103                 'service_ssl_flag': service_ssl_flag,
 104                 'service_type': service_type,
 105             } for i in range(0, count)]
 106         }
 107         self._mock_api_call(api_mock.keep_services().accessible, status, body)
 108         return api_mock
 109
 110     def _mock_api_call(self, mock_method, code, body):
 111         mock_method = mock_method().execute
 112         if code == 200:
 113             mock_method.return_value = body
 114         else:
 115             mock_method.side_effect = arvados.errors.ApiError(
 116                 fake_httplib2_response(code), "{}")
 117
 118
 119 class ArvadosBaseTestCase(unittest.TestCase):
 120     # This class provides common utility functions for our tests.
 121
 122     def setUp(self):
 123         self._tempdirs = []
 124
 125     def tearDown(self):
 126         for workdir in self._tempdirs:
 127             shutil.rmtree(workdir, ignore_errors=True)
 128
 129     def make_tmpdir(self):
 130         self._tempdirs.append(tempfile.mkdtemp())
 131         return self._tempdirs[-1]
 132
 133     def data_file(self, filename):
 134         try:
 135             basedir = os.path.dirname(__file__)
 136         except NameError:
 137             basedir = '.'
 138         return open(os.path.join(basedir, 'data', filename))
 139
 140     def build_directory_tree(self, tree):
 141         tree_root = self.make_tmpdir()
 142         for leaf in tree:
 143             path = os.path.join(tree_root, leaf)
 144             try:
 145                 os.makedirs(os.path.dirname(path))
 146             except OSError as error:
 147                 if error.errno != errno.EEXIST:
 148                     raise
 149             with open(path, 'w') as tmpfile:
 150                 tmpfile.write(leaf)
 151         return tree_root
 152
 153     def make_test_file(self, text="test"):
 154         testfile = tempfile.NamedTemporaryFile()
 155         testfile.write(text)
 156         testfile.flush()
 157         return testfile