diff -r abb05272bc90 Lib/unittest/mock.py --- a/Lib/unittest/mock.py Tue Jan 19 20:14:39 2016 -0500 +++ b/Lib/unittest/mock.py Wed Jan 20 12:45:17 2016 +0900 @@ -25,6 +25,7 @@ import inspect +import asyncio import pprint import sys import builtins @@ -185,9 +186,18 @@ if not name.isidentifier(): name = 'funcopy' context = {'_checksig_': checksig, 'mock': mock} - src = """def %s(*args, **kwargs): + if inspect.iscoroutinefunction(original): + _def, _ret = 'async def', 'await' + elif asyncio.iscoroutinefunction(original): + context['asyncio'] = asyncio + _def, _ret = '@asyncio.coroutine\ndef', 'return' + elif inspect.isgeneratorfunction(original): + _def, _ret = 'def', 'yield' + else: + _def, _ret = 'def', 'return' + src = """%s %s(*args, **kwargs): _checksig_(*args, **kwargs) - return mock(*args, **kwargs)""" % name + %s mock(*args, **kwargs)""" % (_def, name, _ret) exec (src, context) funcopy = context[name] _setup_func(funcopy, mock) @@ -444,6 +454,7 @@ _eat_self=False): _spec_class = None _spec_signature = None + _spec_code = None if spec is not None and not _is_list(spec): if isinstance(spec, type): @@ -453,6 +464,8 @@ res = _get_signature_object(spec, _spec_as_instance, _eat_self) _spec_signature = res and res[1] + if inspect.isfunction(spec) or inspect.ismethod(spec): + _spec_code = spec.__code__ spec = dir(spec) @@ -461,6 +474,7 @@ __dict__['_spec_set'] = spec_set __dict__['_spec_signature'] = _spec_signature __dict__['_mock_methods'] = spec + __dict__['__code__'] = _spec_code def __get_return_value(self): diff -r abb05272bc90 Lib/unittest/test/testmock/testmock.py --- a/Lib/unittest/test/testmock/testmock.py Tue Jan 19 20:14:39 2016 -0500 +++ b/Lib/unittest/test/testmock/testmock.py Wed Jan 20 12:45:17 2016 +0900 @@ -1,6 +1,7 @@ import copy import sys import tempfile +import asyncio import unittest from unittest.test.testmock.support import is_instance @@ -38,6 +39,16 @@ def smeth(a, b, c, d=None): pass + def gmeth(self): + yield 1 + + async def cfmeth(self): + pass + + @asyncio.coroutine + def cometh(self): + pass + class MockTest(unittest.TestCase): @@ -1485,6 +1496,46 @@ self.assertIsInstance(mock, int) mock.foo + def test_isgeneratorfunction(self): + import inspect + def g(): + yield 1 + + self.assertFalse(inspect.isgeneratorfunction(Mock())) + self.assertTrue(inspect.isgeneratorfunction(Mock(g))) + self.assertTrue(inspect.isgeneratorfunction(create_autospec(g))) + self.assertTrue(inspect.isgeneratorfunction(Mock(Something().gmeth))) + self.assertTrue( + inspect.isgeneratorfunction(create_autospec(Something().gmeth))) + + def test_iscoroutinefunction(self): + import inspect + async def f(): + pass + + for m in (inspect, asyncio): + self.assertFalse(m.iscoroutinefunction(Mock())) + self.assertTrue(m.iscoroutinefunction(Mock(f))) + self.assertTrue(m.iscoroutinefunction(create_autospec(f))) + self.assertTrue( + m.iscoroutinefunction(Mock(Something().cfmeth))) + self.assertTrue( + m.iscoroutinefunction(create_autospec(Something().cfmeth))) + + @asyncio.coroutine + def ac(): + pass + # For a function decorated by asyncio.coroutine, + # asyncio.iscoroutinefunction returns True but + # inspect.iscoroutinefunction returns False + self.assertFalse(inspect.iscoroutinefunction(Mock(ac))) + self.assertTrue(asyncio.iscoroutinefunction(Mock(ac))) + self.assertTrue(asyncio.iscoroutinefunction(Mock(Something().cometh))) + self.assertFalse(inspect.iscoroutinefunction(create_autospec(ac))) + self.assertTrue(asyncio.iscoroutinefunction(create_autospec(ac))) + self.assertTrue( + asyncio.iscoroutinefunction(create_autospec(Something().cometh))) + if __name__ == '__main__': unittest.main()