sdk/python/tests/arvados_testutil.py

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