sdk/python/tests/arvados_testutil.py

   1 #!/usr/bin/env python
   2
   3 import errno
   4 import os
   5 import shutil
   6 import tempfile
   7 import unittest
   8
   9 class ArvadosBaseTestCase(unittest.TestCase):
  10     # This class provides common utility functions for our tests.
  11
  12     def setUp(self):
  13         self._tempdirs = []
  14
  15     def tearDown(self):
  16         for workdir in self._tempdirs:
  17             shutil.rmtree(workdir, ignore_errors=True)
  18
  19     def make_tmpdir(self):
  20         self._tempdirs.append(tempfile.mkdtemp())
  21         return self._tempdirs[-1]
  22
  23     def data_file(self, filename):
  24         try:
  25             basedir = os.path.dirname(__file__)
  26         except NameError:
  27             basedir = '.'
  28         return open(os.path.join(basedir, 'data', filename))
  29
  30
  31 class ArvadosKeepLocalStoreTestCase(ArvadosBaseTestCase):
  32     def setUp(self):
  33         super(ArvadosKeepLocalStoreTestCase, self).setUp()
  34         self._orig_keep_local_store = os.environ.get('KEEP_LOCAL_STORE')
  35         os.environ['KEEP_LOCAL_STORE'] = self.make_tmpdir()
  36
  37     def tearDown(self):
  38         if self._orig_keep_local_store is None:
  39             del os.environ['KEEP_LOCAL_STORE']
  40         else:
  41             os.environ['KEEP_LOCAL_STORE'] = self._orig_keep_local_store
  42         super(ArvadosKeepLocalStoreTestCase, self).tearDown()
  43
  44     def build_directory_tree(self, tree):
  45         tree_root = self.make_tmpdir()
  46         for leaf in tree:
  47             path = os.path.join(tree_root, leaf)
  48             try:
  49                 os.makedirs(os.path.dirname(path))
  50             except OSError as error:
  51                 if error.errno != errno.EEXIST:
  52                     raise
  53             with open(path, 'w') as tmpfile:
  54                 tmpfile.write(leaf)
  55         return tree_root
  56
  57     def make_test_file(self, text="test"):
  58         testfile = tempfile.NamedTemporaryFile()
  59         testfile.write(text)
  60         testfile.flush()
  61         return testfile