Adding patch.stopall method to unittest.mock

This commit is contained in:
Michael Foord 2012-06-10 20:36:32 +01:00
parent bfcb42936b
commit f7c4158057
3 changed files with 44 additions and 4 deletions

View File

@ -1354,8 +1354,12 @@ method of a `TestCase`:
As an added bonus you no longer need to keep a reference to the `patcher`
object.
In fact `start` and `stop` are just aliases for the context manager
`__enter__` and `__exit__` methods.
It is also possible to stop all patches which have been started by using
`patch.stopall`.
.. function:: patch.stopall
Stop all active patches.
TEST_PREFIX

View File

@ -1002,6 +1002,7 @@ def _is_started(patcher):
class _patch(object):
attribute_name = None
_active_patches = set()
def __init__(
self, getter, attribute, new, spec, create,
@ -1270,8 +1271,18 @@ class _patch(object):
if _is_started(patcher):
patcher.__exit__(*exc_info)
start = __enter__
stop = __exit__
def start(self):
"""Activate a patch, returning any created mock."""
result = self.__enter__()
self._active_patches.add(self)
return result
def stop(self):
"""Stop an active patch."""
self._active_patches.discard(self)
return self.__exit__()
@ -1562,9 +1573,16 @@ def _clear_dict(in_dict):
del in_dict[key]
def _patch_stopall():
"""Stop all active patches."""
for patch in list(_patch._active_patches):
patch.stop()
patch.object = _patch_object
patch.dict = _patch_dict
patch.multiple = _patch_multiple
patch.stopall = _patch_stopall
patch.TEST_PREFIX = 'test'
magic_methods = (

View File

@ -1762,6 +1762,24 @@ class PatchTest(unittest.TestCase):
p.stop()
def test_patch_stopall(self):
unlink = os.unlink
chdir = os.chdir
path = os.path
patch('os.unlink', something).start()
patch('os.chdir', something_else).start()
@patch('os.path')
def patched(mock_path):
patch.stopall()
self.assertIs(os.path, mock_path)
self.assertIs(os.unlink, unlink)
self.assertIs(os.chdir, chdir)
patched()
self.assertIs(os.path, path)
if __name__ == '__main__':
unittest.main()