sdk/python/tests/arvados_testutil.py

   1 #!/usr/bin/env python
   2
   3 import errno
   4 import httplib
   5 import httplib2
   6 import mock
   7 import os
   8 import shutil
   9 import tempfile
  10 import unittest
  11
  12 # Use this hostname when you want to make sure the traffic will be
  13 # instantly refused.  100::/64 is a dedicated black hole.
  14 TEST_HOST = '100::'
  15
  16 skip_sleep = mock.patch('time.sleep', lambda n: None)  # clown'll eat me
  17
  18 def fake_httplib2_response(code, **headers):
  19     headers.update(status=str(code),
  20                    reason=httplib.responses.get(code, "Unknown Response"))
  21     return httplib2.Response(headers)
  22
  23 def mock_responses(body, *codes, **headers):
  24     return mock.patch('httplib2.Http.request', side_effect=(
  25             (fake_httplib2_response(code, **headers), body) for code in codes))
  26
  27 class ArvadosBaseTestCase(unittest.TestCase):
  28     # This class provides common utility functions for our tests.
  29
  30     def setUp(self):
  31         self._tempdirs = []
  32
  33     def tearDown(self):
  34         for workdir in self._tempdirs:
  35             shutil.rmtree(workdir, ignore_errors=True)
  36
  37     def make_tmpdir(self):
  38         self._tempdirs.append(tempfile.mkdtemp())
  39         return self._tempdirs[-1]
  40
  41     def data_file(self, filename):
  42         try:
  43             basedir = os.path.dirname(__file__)
  44         except NameError:
  45             basedir = '.'
  46         return open(os.path.join(basedir, 'data', filename))
  47
  48     def build_directory_tree(self, tree):
  49         tree_root = self.make_tmpdir()
  50         for leaf in tree:
  51             path = os.path.join(tree_root, leaf)
  52             try:
  53                 os.makedirs(os.path.dirname(path))
  54             except OSError as error:
  55                 if error.errno != errno.EEXIST:
  56                     raise
  57             with open(path, 'w') as tmpfile:
  58                 tmpfile.write(leaf)
  59         return tree_root
  60
  61     def make_test_file(self, text="test"):
  62         testfile = tempfile.NamedTemporaryFile()
  63         testfile.write(text)
  64         testfile.flush()
  65         return testfile