| 1 | import shutil
|
|---|
| 2 | import tempfile
|
|---|
| 3 | import unittest
|
|---|
| 4 |
|
|---|
| 5 |
|
|---|
| 6 | class SetupStack(object):
|
|---|
| 7 |
|
|---|
| 8 | def __init__(self):
|
|---|
| 9 | self._on_teardown = []
|
|---|
| 10 |
|
|---|
| 11 | def add_teardown(self, teardown):
|
|---|
| 12 | self._on_teardown.append(teardown)
|
|---|
| 13 |
|
|---|
| 14 | def tear_down(self):
|
|---|
| 15 | for func in reversed(self._on_teardown):
|
|---|
| 16 | func()
|
|---|
| 17 |
|
|---|
| 18 |
|
|---|
| 19 | class TearDownConvenience(object):
|
|---|
| 20 |
|
|---|
| 21 | def __init__(self, setup_stack=None):
|
|---|
| 22 | self._own_setup_stack = setup_stack is None
|
|---|
| 23 | if setup_stack is None:
|
|---|
| 24 | setup_stack = SetupStack()
|
|---|
| 25 | self._setup_stack = setup_stack
|
|---|
| 26 |
|
|---|
| 27 | # only call this convenience method if no setup_stack was supplied to c'tor
|
|---|
| 28 | def tear_down(self):
|
|---|
| 29 | assert self._own_setup_stack
|
|---|
| 30 | self._setup_stack.tear_down()
|
|---|
| 31 |
|
|---|
| 32 |
|
|---|
| 33 | class TempDirMaker(TearDownConvenience):
|
|---|
| 34 |
|
|---|
| 35 | def make_temp_dir(self):
|
|---|
| 36 | temp_dir = tempfile.mkdtemp(prefix="tmp-%s-" % self.__class__.__name__)
|
|---|
| 37 | def tear_down():
|
|---|
| 38 | shutil.rmtree(temp_dir)
|
|---|
| 39 | self._setup_stack.add_teardown(tear_down)
|
|---|
| 40 | return temp_dir
|
|---|
| 41 |
|
|---|
| 42 |
|
|---|
| 43 | class MonkeyPatcher(TearDownConvenience):
|
|---|
| 44 |
|
|---|
| 45 | def monkey_patch(self, obj, name, value):
|
|---|
| 46 | orig_value = getattr(obj, name)
|
|---|
| 47 | setattr(obj, name, value)
|
|---|
| 48 | def reverse_patch():
|
|---|
| 49 | setattr(obj, name, orig_value)
|
|---|
| 50 | self._setup_stack.add_teardown(reverse_patch)
|
|---|
| 51 |
|
|---|
| 52 |
|
|---|
| 53 | class TestCase(unittest.TestCase):
|
|---|
| 54 |
|
|---|
| 55 | def setUp(self):
|
|---|
| 56 | self._setup_stack = SetupStack()
|
|---|
| 57 |
|
|---|
| 58 | def tearDown(self):
|
|---|
| 59 | self._setup_stack.tear_down()
|
|---|
| 60 |
|
|---|
| 61 | def make_temp_dir(self, *args, **kwds):
|
|---|
| 62 | return TempDirMaker(self._setup_stack).make_temp_dir(*args, **kwds)
|
|---|
| 63 |
|
|---|
| 64 | def monkey_patch(self, *args, **kwds):
|
|---|
| 65 | return MonkeyPatcher(self._setup_stack).monkey_patch(*args, **kwds)
|
|---|
| 66 |
|
|---|
| 67 | def assert_contains(self, container, containee):
|
|---|
| 68 | self.assertTrue(containee in container, "%r not in %r" %
|
|---|
| 69 | (containee, container))
|
|---|
| 70 |
|
|---|
| 71 | def assert_less_than(self, got, expected):
|
|---|
| 72 | self.assertTrue(got < expected, "%r >= %r" %
|
|---|
| 73 | (got, expected))
|
|---|
