Skip to content

Commit

Permalink
pythongh-98086: Now patch.dict can decorate async functions
Browse files Browse the repository at this point in the history
  • Loading branch information
sobolevn committed Oct 8, 2022
1 parent 4ed00be commit ecb48f1
Show file tree
Hide file tree
Showing 3 changed files with 36 additions and 0 deletions.
17 changes: 17 additions & 0 deletions Lib/test/test_unittest/testmock/testasync.py
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,23 @@ async def test_async():

run(test_async())

def test_patch_dict_async_def(self):
foo = {'a': 'a'}
@patch.dict(foo, {'a': 'b'})
async def test_async():
self.assertEqual(foo['a'], 'b')

self.assertTrue(iscoroutinefunction(test_async))
run(test_async())

def test_patch_dict_async_def_context(self):
foo = {'a': 'a'}
async def test_async():
with patch.dict(foo, {'a': 'b'}):
self.assertEqual(foo['a'], 'b')

run(test_async())


class AsyncMockTest(unittest.TestCase):
def test_iscoroutinefunction_default(self):
Expand Down
18 changes: 18 additions & 0 deletions Lib/unittest/mock.py
Original file line number Diff line number Diff line change
Expand Up @@ -1799,6 +1799,12 @@ def __init__(self, in_dict, values=(), clear=False, **kwargs):
def __call__(self, f):
if isinstance(f, type):
return self.decorate_class(f)
if inspect.iscoroutinefunction(f):
return self.decorate_async_callable(f)
return self.decorate_callable(f)


def decorate_callable(self, f):
@wraps(f)
def _inner(*args, **kw):
self._patch_dict()
Expand All @@ -1810,6 +1816,18 @@ def _inner(*args, **kw):
return _inner


def decorate_async_callable(self, f):
@wraps(f)
async def _inner(*args, **kw):
self._patch_dict()
try:
return await f(*args, **kw)
finally:
self._unpatch_dict()

return _inner


def decorate_class(self, klass):
for attr in dir(klass):
attr_value = getattr(klass, attr)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Make sure ``patch.dict()`` can be applied on async functions.

0 comments on commit ecb48f1

Please sign in to comment.