Skip to content

Commit f78011f

Browse files
committed
Limit starting a patcher more than once without stopping it
Previously, this would cause an `AttributeError` if the patch stopped more than once after this, and would also disrupt the original patched object.
1 parent 0f6bb28 commit f78011f

File tree

3 files changed

+81
-17
lines changed

3 files changed

+81
-17
lines changed

Lib/test/test_unittest/testmock/testpatch.py

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -745,6 +745,35 @@ def test_stop_idempotent(self):
745745
self.assertIsNone(patcher.stop())
746746

747747

748+
def test_exit_idempotent(self):
749+
patcher = patch(foo_name, 'bar', 3)
750+
with patcher:
751+
patcher.stop()
752+
753+
754+
def test_second_start_failure(self):
755+
patcher = patch(foo_name, 'bar', 3)
756+
patcher.start()
757+
try:
758+
self.assertRaises(RuntimeError, patcher.start)
759+
finally:
760+
patcher.stop()
761+
762+
763+
def test_second_enter_failure(self):
764+
patcher = patch(foo_name, 'bar', 3)
765+
with patcher:
766+
self.assertRaises(RuntimeError, patcher.start)
767+
768+
769+
def test_second_start_after_stop(self):
770+
patcher = patch(foo_name, 'bar', 3)
771+
patcher.start()
772+
patcher.stop()
773+
patcher.start()
774+
patcher.stop()
775+
776+
748777
def test_patchobject_start_stop(self):
749778
original = something
750779
patcher = patch.object(PTModule, 'something', 'foo')
@@ -1098,7 +1127,7 @@ def test_new_callable_patch(self):
10981127

10991128
self.assertIsNot(m1, m2)
11001129
for mock in m1, m2:
1101-
self.assertNotCallable(m1)
1130+
self.assertNotCallable(mock)
11021131

11031132

11041133
def test_new_callable_patch_object(self):
@@ -1111,7 +1140,7 @@ def test_new_callable_patch_object(self):
11111140

11121141
self.assertIsNot(m1, m2)
11131142
for mock in m1, m2:
1114-
self.assertNotCallable(m1)
1143+
self.assertNotCallable(mock)
11151144

11161145

11171146
def test_new_callable_keyword_arguments(self):

Lib/unittest/mock.py

Lines changed: 45 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1320,6 +1320,14 @@ def _check_spec_arg_typos(kwargs_to_check):
13201320
)
13211321

13221322

1323+
class _PatchStartedContext(object):
1324+
def __init__(self, exit_stack, is_local, target, temp_original):
1325+
self.exit_stack = exit_stack
1326+
self.is_local = is_local
1327+
self.target = target
1328+
self.temp_original = temp_original
1329+
1330+
13231331
class _patch(object):
13241332

13251333
attribute_name = None
@@ -1360,6 +1368,7 @@ def __init__(
13601368
self.autospec = autospec
13611369
self.kwargs = kwargs
13621370
self.additional_patchers = []
1371+
self._started_context = None
13631372

13641373

13651374
def copy(self):
@@ -1469,13 +1478,31 @@ def get_original(self):
14691478
)
14701479
return original, local
14711480

1481+
@property
1482+
def is_started(self):
1483+
return self._started_context is not None
1484+
1485+
@property
1486+
def is_local(self):
1487+
return self._started_context.is_local
1488+
1489+
@property
1490+
def target(self):
1491+
return self._started_context.target
1492+
1493+
@property
1494+
def temp_original(self):
1495+
return self._started_context.temp_original
14721496

14731497
def __enter__(self):
14741498
"""Perform the patch."""
1499+
if self.is_started:
1500+
raise RuntimeError("Patch is already started")
1501+
14751502
new, spec, spec_set = self.new, self.spec, self.spec_set
14761503
autospec, kwargs = self.autospec, self.kwargs
14771504
new_callable = self.new_callable
1478-
self.target = self.getter()
1505+
target = self.getter()
14791506

14801507
# normalise False to None
14811508
if spec is False:
@@ -1491,7 +1518,7 @@ def __enter__(self):
14911518
spec_set not in (True, None)):
14921519
raise TypeError("Can't provide explicit spec_set *and* spec or autospec")
14931520

1494-
original, local = self.get_original()
1521+
original, is_local = self.get_original()
14951522

14961523
if new is DEFAULT and autospec is None:
14971524
inherit = False
@@ -1579,17 +1606,17 @@ def __enter__(self):
15791606
if autospec is True:
15801607
autospec = original
15811608

1582-
if _is_instance_mock(self.target):
1609+
if _is_instance_mock(target):
15831610
raise InvalidSpecError(
15841611
f'Cannot autospec attr {self.attribute!r} as the patch '
15851612
f'target has already been mocked out. '
1586-
f'[target={self.target!r}, attr={autospec!r}]')
1613+
f'[target={target!r}, attr={autospec!r}]')
15871614
if _is_instance_mock(autospec):
1588-
target_name = getattr(self.target, '__name__', self.target)
1615+
target_name = getattr(target, '__name__', target)
15891616
raise InvalidSpecError(
15901617
f'Cannot autospec attr {self.attribute!r} from target '
15911618
f'{target_name!r} as it has already been mocked out. '
1592-
f'[target={self.target!r}, attr={autospec!r}]')
1619+
f'[target={target!r}, attr={autospec!r}]')
15931620

15941621
new = create_autospec(autospec, spec_set=spec_set,
15951622
_name=self.attribute, **kwargs)
@@ -1600,17 +1627,20 @@ def __enter__(self):
16001627

16011628
new_attr = new
16021629

1603-
self.temp_original = original
1604-
self.is_local = local
1605-
self._exit_stack = contextlib.ExitStack()
1630+
self._started_context = _PatchStartedContext(
1631+
exit_stack=contextlib.ExitStack(),
1632+
is_local=is_local,
1633+
target=self.getter(),
1634+
temp_original=original,
1635+
)
16061636
try:
16071637
setattr(self.target, self.attribute, new_attr)
16081638
if self.attribute_name is not None:
16091639
extra_args = {}
16101640
if self.new is DEFAULT:
16111641
extra_args[self.attribute_name] = new
16121642
for patching in self.additional_patchers:
1613-
arg = self._exit_stack.enter_context(patching)
1643+
arg = self._started_context.exit_stack.enter_context(patching)
16141644
if patching.new is DEFAULT:
16151645
extra_args.update(arg)
16161646
return extra_args
@@ -1622,6 +1652,9 @@ def __enter__(self):
16221652

16231653
def __exit__(self, *exc_info):
16241654
"""Undo the patch."""
1655+
if not self.is_started:
1656+
return
1657+
16251658
if self.is_local and self.temp_original is not DEFAULT:
16261659
setattr(self.target, self.attribute, self.temp_original)
16271660
else:
@@ -1633,11 +1666,8 @@ def __exit__(self, *exc_info):
16331666
# needed for proxy objects like django settings
16341667
setattr(self.target, self.attribute, self.temp_original)
16351668

1636-
del self.temp_original
1637-
del self.is_local
1638-
del self.target
1639-
exit_stack = self._exit_stack
1640-
del self._exit_stack
1669+
exit_stack = self._started_context.exit_stack
1670+
self._started_context = None
16411671
return exit_stack.__exit__(*exc_info)
16421672

16431673

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
Limit starting a patcher (from :func:`unittest.mock.patch`,
2+
:func:`unittest.mock.patch.object` or :func:`unittest.mock.dict`) more than
3+
once without stopping it. Previously, this would cause an
4+
:exc:`AttributeError` if the patch stopped more than once after this, and
5+
would also disrupt the original patched object.

0 commit comments

Comments
 (0)