# -*- test-case-name: twisted.test.test_monkey -*- # Copyright (c) Twisted Matrix Laboratories. # See LICENSE for details. from __future__ import division, absolute_import class MonkeyPatcher(object): """ Cover up attributes with new objects. Neat for monkey-patching things for unit-testing purposes. """ def __init__(self, *patches): # List of patches to apply in (obj, name, value). self._patchesToApply = [] # List of the original values for things that have been patched. # (obj, name, value) format. self._originals = [] for patch in patches: self.addPatch(*patch) def addPatch(self, obj, name, value): """ Add a patch so that the attribute C{name} on C{obj} will be assigned to C{value} when C{patch} is called or during C{runWithPatches}. You can restore the original values with a call to restore(). """ self._patchesToApply.append((obj, name, value)) def _alreadyPatched(self, obj, name): """ Has the C{name} attribute of C{obj} already been patched by this patcher? """ for o, n, v in self._originals: if (o, n) == (obj, name): return True return False def patch(self): """ Apply all of the patches that have been specified with L{addPatch}. Reverse this operation using L{restore}. """ for obj, name, value in self._patchesToApply: if not self._alreadyPatched(obj, name): self._originals.append((obj, name, getattr(obj, name))) setattr(obj, name, value) def restore(self): """ Restore all original values to any patched objects. """ while self._originals: obj, name, value = self._originals.pop() setattr(obj, name, value) def runWithPatches(self, f, *args, **kw): """ Apply each patch already specified. Then run the function f with the given args and kwargs. Restore everything when done. """ self.patch() try: return f(*args, **kw) finally: self.restore()