X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/0199da877b222c751a31d70a6617b92c2a56fe70..2693a197ea37490a891c54bbb279937a7b4e897c:/sdk/python/tests/arvados_testutil.py diff --git a/sdk/python/tests/arvados_testutil.py b/sdk/python/tests/arvados_testutil.py index 5992371902..0035659796 100644 --- a/sdk/python/tests/arvados_testutil.py +++ b/sdk/python/tests/arvados_testutil.py @@ -1,4 +1,6 @@ -#!/usr/bin/env python +# Copyright (C) The Arvados Authors. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 from future import standard_library standard_library.install_aliases() @@ -21,6 +23,12 @@ import sys import tempfile import unittest +if sys.version_info >= (3, 0): + from io import StringIO, BytesIO +else: + from cStringIO import StringIO + BytesIO = StringIO + # Use this hostname when you want to make sure the traffic will be # instantly refused. 100::/64 is a dedicated black hole. TEST_HOST = '100::' @@ -47,33 +55,55 @@ def fake_httplib2_response(code, **headers): return httplib2.Response(headers) def mock_responses(body, *codes, **headers): + if not isinstance(body, bytes) and hasattr(body, 'encode'): + body = body.encode() return mock.patch('httplib2.Http.request', side_effect=queue_with(( (fake_httplib2_response(code, **headers), body) for code in codes))) def mock_api_responses(api_client, body, codes, headers={}): + if not isinstance(body, bytes) and hasattr(body, 'encode'): + body = body.encode() return mock.patch.object(api_client._http, 'request', side_effect=queue_with(( (fake_httplib2_response(code, **headers), body) for code in codes))) def str_keep_locator(s): - return '{}+{}'.format(hashlib.md5(s).hexdigest(), len(s)) + return '{}+{}'.format(hashlib.md5(s if isinstance(s, bytes) else s.encode()).hexdigest(), len(s)) @contextlib.contextmanager def redirected_streams(stdout=None, stderr=None): + if stdout == StringIO: + stdout = StringIO() + if stderr == StringIO: + stderr = StringIO() orig_stdout, sys.stdout = sys.stdout, stdout or sys.stdout orig_stderr, sys.stderr = sys.stderr, stderr or sys.stderr try: - yield + yield (stdout, stderr) finally: sys.stdout = orig_stdout sys.stderr = orig_stderr +class VersionChecker(object): + def assertVersionOutput(self, out, err): + if sys.version_info >= (3, 0): + self.assertEqual(err.getvalue(), '') + v = out.getvalue() + else: + # Python 2 writes version info on stderr. + self.assertEqual(out.getvalue(), '') + v = err.getvalue() + self.assertRegex(v, r"[0-9]+\.[0-9]+\.[0-9]+(\.dev[0-9]+)?$\n") + + class FakeCurl(object): @classmethod - def make(cls, code, body='', headers={}): + def make(cls, code, body=b'', headers={}): + if not isinstance(body, bytes) and hasattr(body, 'encode'): + body = body.encode() return mock.Mock(spec=cls, wraps=cls(code, body, headers)) - def __init__(self, code=200, body='', headers={}): + def __init__(self, code=200, body=b'', headers={}): self._opt = {} self._got_url = None self._writer = None @@ -146,7 +176,9 @@ def mock_keep_responses(body, *codes, **headers): class MockStreamReader(object): def __init__(self, name='.', *data): self._name = name - self._data = ''.join(data) + self._data = b''.join([ + b if isinstance(b, bytes) else b.encode() + for b in data]) self._data_locators = [str_keep_locator(d) for d in data] self.num_retries = 0 @@ -158,7 +190,13 @@ class MockStreamReader(object): class ApiClientMock(object): def api_client_mock(self): - return mock.MagicMock(name='api_client_mock') + api_mock = mock.MagicMock(name='api_client_mock') + api_mock.config.return_value = { + 'StorageClasses': { + 'default': {'Default': True} + } + } + return api_mock def mock_keep_services(self, api_mock=None, status=200, count=12, service_type='disk', @@ -190,7 +228,7 @@ class ApiClientMock(object): mock_method.return_value = body else: mock_method.side_effect = arvados.errors.ApiError( - fake_httplib2_response(code), "{}") + fake_httplib2_response(code), b"{}") class ArvadosBaseTestCase(unittest.TestCase): @@ -227,8 +265,45 @@ class ArvadosBaseTestCase(unittest.TestCase): tmpfile.write(leaf) return tree_root - def make_test_file(self, text="test"): + def make_test_file(self, text=b"test"): testfile = tempfile.NamedTemporaryFile() testfile.write(text) testfile.flush() return testfile + +if sys.version_info < (3, 0): + # There is no assert[Not]Regex that works in both Python 2 and 3, + # so we backport Python 3 style to Python 2. + def assertRegex(self, *args, **kwargs): + return self.assertRegexpMatches(*args, **kwargs) + def assertNotRegex(self, *args, **kwargs): + return self.assertNotRegexpMatches(*args, **kwargs) + unittest.TestCase.assertRegex = assertRegex + unittest.TestCase.assertNotRegex = assertNotRegex + +def binary_compare(a, b): + if len(a) != len(b): + return False + for i in range(0, len(a)): + if a[i] != b[i]: + return False + return True + +def make_block_cache(disk_cache): + if disk_cache: + disk_cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "arvados", "keep") + shutil.rmtree(disk_cache_dir, ignore_errors=True) + block_cache = arvados.keep.KeepBlockCache(disk_cache=disk_cache) + return block_cache + + +class DiskCacheBase: + def make_block_cache(self, disk_cache): + self.disk_cache_dir = tempfile.mkdtemp() if disk_cache else None + block_cache = arvados.keep.KeepBlockCache(disk_cache=disk_cache, + disk_cache_dir=self.disk_cache_dir) + return block_cache + + def tearDown(self): + if self.disk_cache_dir: + shutil.rmtree(self.disk_cache_dir)