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_get_responses(body, *codes, **headers):
  43     return mock.patch('requests.get', side_effect=(
  44         fake_requests_response(code, body, **headers) for code in codes))
  45
  46 def mock_put_responses(body, *codes, **headers):
  47     return mock.patch('requests.put', side_effect=(
  48         fake_requests_response(code, body, **headers) for code in codes))
  49
  50 def mock_requestslib_responses(method, body, *codes, **headers):
  51     return mock.patch(method, side_effect=(
  52         fake_requests_response(code, body, **headers) for code in codes))
  53
  54 class MockStreamReader(object):
  55     def __init__(self, name='.', *data):
  56         self._name = name
  57         self._data = ''.join(data)
  58         self._data_locators = ['{}+{}'.format(hashlib.md5(d).hexdigest(),
  59                                               len(d)) for d in data]
  60         self.num_retries = 0
  61
  62     def name(self):
  63         return self._name
  64
  65     def readfrom(self, start, size, num_retries=None):
  66         return self._data[start:start + size]
  67
  68
  69 class ArvadosBaseTestCase(unittest.TestCase):
  70     # This class provides common utility functions for our tests.
  71
  72     def setUp(self):
  73         self._tempdirs = []
  74
  75     def tearDown(self):
  76         for workdir in self._tempdirs:
  77             shutil.rmtree(workdir, ignore_errors=True)
  78
  79     def make_tmpdir(self):
  80         self._tempdirs.append(tempfile.mkdtemp())
  81         return self._tempdirs[-1]
  82
  83     def data_file(self, filename):
  84         try:
  85             basedir = os.path.dirname(__file__)
  86         except NameError:
  87             basedir = '.'
  88         return open(os.path.join(basedir, 'data', filename))
  89
  90     def build_directory_tree(self, tree):
  91         tree_root = self.make_tmpdir()
  92         for leaf in tree:
  93             path = os.path.join(tree_root, leaf)
  94             try:
  95                 os.makedirs(os.path.dirname(path))
  96             except OSError as error:
  97                 if error.errno != errno.EEXIST:
  98                     raise
  99             with open(path, 'w') as tmpfile:
 100                 tmpfile.write(leaf)
 101         return tree_root
 102
 103     def make_test_file(self, text="test"):
 104         testfile = tempfile.NamedTemporaryFile()
 105         testfile.write(text)
 106         testfile.flush()
 107         return testfile