Add a convenience function to just patch a single thing.

This commit is contained in:
Jonathan Lange
2010-08-04 14:01:36 +01:00
parent a102e886c5
commit 77b1d999e1
2 changed files with 41 additions and 1 deletions

View File

@@ -2,6 +2,11 @@
"""Helpers for monkey-patching Python code."""
__all__ = [
'MonkeyPatcher',
'patch',
]
class MonkeyPatcher(object):
"""A set of monkey-patches that can be applied and removed all together.
@@ -73,3 +78,20 @@ class MonkeyPatcher(object):
return f(*args, **kw)
finally:
self.restore()
def patch(obj, attribute, value):
"""Set 'obj.attribute' to 'value' and return a callable to restore 'obj'.
If 'attribute' is not set on 'obj' already, then the returned callable
will delete the attribute when called.
:param obj: An object to monkey-patch.
:param attribute: The name of the attribute to patch.
:param value: The value to set 'obj.attribute' to.
:return: A nullary callable that, when run, will restore 'obj' to its
original state.
"""
patcher = MonkeyPatcher((obj, attribute, value))
patcher.patch()
return patcher.restore

View File

@@ -4,7 +4,7 @@
"""Tests for testtools.monkey."""
from testtools import TestCase
from testtools.monkey import MonkeyPatcher
from testtools.monkey import MonkeyPatcher, patch
class TestObj:
@@ -143,6 +143,24 @@ class MonkeyPatcherTest(TestCase):
self.assertEquals(self.test_object.bar, self.original_object.bar)
class TestPatchHelper(TestCase):
def test_patch_patches(self):
# patch(obj, name, value) sets obj.name to value.
test_object = TestObj()
patch(test_object, 'foo', 42)
self.assertEqual(42, test_object.foo)
def test_patch_returns_cleanup(self):
# patch(obj, name, value) returns a nullary callable that restores obj
# to its original state when run.
test_object = TestObj()
original = test_object.foo
cleanup = patch(test_object, 'foo', 42)
cleanup()
self.assertEqual(original, test_object.foo)
def test_suite():
from unittest import TestLoader
return TestLoader().loadTestsFromName(__name__)