Merge branch 'master' into 4904-arv-web
[arvados.git] / 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