@@ -377,43 +377,88 @@ def test_add_side_effect_iterable(self):
377377
378378class AsyncContextManagerTest (unittest .TestCase ):
379379 class WithAsyncContextManager :
380- def __init__ (self ):
381- self .entered = False
382- self .exited = False
383-
384380 async def __aenter__ (self , * args , ** kwargs ):
385381 self .entered = True
386382 return self
387383
388384 async def __aexit__ (self , * args , ** kwargs ):
389385 self .exited = True
390386
391- def test_magic_methods_are_async_mocks (self ):
392- mock = MagicMock (self .WithAsyncContextManager ())
393- self .assertIsInstance (mock .__aenter__ , AsyncMock )
394- self .assertIsInstance (mock .__aexit__ , AsyncMock )
387+ class WithSyncContextManager :
388+ def __enter__ (self , * args , ** kwargs ):
389+ return self
390+
391+ def __exit__ (self , * args , ** kwargs ):
392+ pass
393+
394+ class ProductionCode :
395+ # Example real-world(ish) code
396+ def __init__ (self ):
397+ self .session = None
398+
399+ async def main (self ):
400+ async with self .session .post ('https://python.org' ) as response :
401+ val = await response .json ()
402+ return val
403+
404+ def test_async_magic_methods_are_async_mocks_with_magicmock (self ):
405+ cm_mock = MagicMock (self .WithAsyncContextManager ())
406+ self .assertIsInstance (cm_mock .__aenter__ , AsyncMock )
407+ self .assertIsInstance (cm_mock .__aexit__ , AsyncMock )
408+
409+ def test_magicmock_has_async_magic_methods (self ):
410+ cm = MagicMock (name = 'magic_cm' )
411+ self .assertTrue (hasattr (cm , "__aenter__" ))
412+ self .assertTrue (hasattr (cm , "__aexit__" ))
413+
414+ def test_magic_methods_are_async_functions (self ):
415+ cm = MagicMock (name = 'magic_cm' )
416+ self .assertIsInstance (cm .__aenter__ , AsyncMock )
417+ self .assertIsInstance (cm .__aexit__ , AsyncMock )
418+ # AsyncMocks are also coroutine functions
419+ self .assertTrue (asyncio .iscoroutinefunction (cm .__aenter__ ))
420+ self .assertTrue (asyncio .iscoroutinefunction (cm .__aexit__ ))
421+
422+ def test_set_return_value_of_aenter (self ):
423+ def inner_test (mock_type ):
424+ pc = self .ProductionCode ()
425+ pc .session = MagicMock (name = 'sessionmock' )
426+ cm = mock_type (name = 'magic_cm' )
427+ response = AsyncMock (name = 'response' )
428+ response .json = AsyncMock (return_value = {'json' : 123 })
429+ cm .__aenter__ .return_value = response
430+ pc .session .post .return_value = cm
431+ result = asyncio .run (pc .main ())
432+ self .assertEqual (result , {'json' : 123 })
433+
434+ for mock_type in [AsyncMock , MagicMock ]:
435+ with self .subTest (f"test set return value of aenter with { mock_type } " ):
436+ inner_test (mock_type )
395437
396438 def test_mock_supports_async_context_manager (self ):
397- called = False
398- instance = self .WithAsyncContextManager ()
399- mock_instance = MagicMock (instance )
439+ def inner_test (mock_type ):
440+ called = False
441+ cm = self .WithAsyncContextManager ()
442+ cm_mock = mock_type (cm )
443+
444+ async def use_context_manager ():
445+ nonlocal called
446+ async with cm_mock as result :
447+ called = True
448+ return result
400449
401- async def use_context_manager ():
402- nonlocal called
403- async with mock_instance as result :
404- called = True
405- return result
406-
407- result = asyncio .run (use_context_manager ())
408- self .assertFalse (instance .entered )
409- self .assertFalse (instance .exited )
410- self .assertTrue (called )
411- self .assertTrue (mock_instance .entered )
412- self .assertTrue (mock_instance .exited )
413- self .assertTrue (mock_instance .__aenter__ .called )
414- self .assertTrue (mock_instance .__aexit__ .called )
415- self .assertIsNot (mock_instance , result )
416- self .assertIsInstance (result , AsyncMock )
450+ cm_result = asyncio .run (use_context_manager ())
451+ self .assertTrue (called )
452+ self .assertTrue (cm_mock .__aenter__ .called )
453+ self .assertTrue (cm_mock .__aexit__ .called )
454+ cm_mock .__aenter__ .assert_awaited ()
455+ cm_mock .__aexit__ .assert_awaited ()
456+ # We mock __aenter__ so it does not return self
457+ self .assertIsNot (cm_mock , cm_result )
458+
459+ for mock_type in [AsyncMock , MagicMock ]:
460+ with self .subTest (f"test context manager magics with { mock_type } " ):
461+ inner_test (mock_type )
417462
418463 def test_mock_customize_async_context_manager (self ):
419464 instance = self .WithAsyncContextManager ()
@@ -481,27 +526,30 @@ async def __anext__(self):
481526
482527 raise StopAsyncIteration
483528
484- def test_mock_aiter_and_anext (self ):
485- instance = self .WithAsyncIterator ()
486- mock_instance = MagicMock (instance )
487-
488- self .assertEqual (asyncio .iscoroutine (instance .__aiter__ ),
489- asyncio .iscoroutine (mock_instance .__aiter__ ))
490- self .assertEqual (asyncio .iscoroutine (instance .__anext__ ),
491- asyncio .iscoroutine (mock_instance .__anext__ ))
492-
493- iterator = instance .__aiter__ ()
494- if asyncio .iscoroutine (iterator ):
495- iterator = asyncio .run (iterator )
496-
497- mock_iterator = mock_instance .__aiter__ ()
498- if asyncio .iscoroutine (mock_iterator ):
499- mock_iterator = asyncio .run (mock_iterator )
529+ def test_aiter_set_return_value (self ):
530+ mock_iter = AsyncMock (name = "tester" )
531+ mock_iter .__aiter__ .return_value = [1 , 2 , 3 ]
532+ async def main ():
533+ return [i async for i in mock_iter ]
534+ result = asyncio .run (main ())
535+ self .assertEqual (result , [1 , 2 , 3 ])
536+
537+ def test_mock_aiter_and_anext_asyncmock (self ):
538+ def inner_test (mock_type ):
539+ instance = self .WithAsyncIterator ()
540+ mock_instance = mock_type (instance )
541+ # Check that the mock and the real thing bahave the same
542+ # __aiter__ is not actually async, so not a coroutinefunction
543+ self .assertFalse (asyncio .iscoroutinefunction (instance .__aiter__ ))
544+ self .assertFalse (asyncio .iscoroutinefunction (mock_instance .__aiter__ ))
545+ # __anext__ is async
546+ self .assertTrue (asyncio .iscoroutinefunction (instance .__anext__ ))
547+ self .assertTrue (asyncio .iscoroutinefunction (mock_instance .__anext__ ))
548+
549+ for mock_type in [AsyncMock , MagicMock ]:
550+ with self .subTest (f"test aiter and anext corourtine with { mock_type } " ):
551+ inner_test (mock_type )
500552
501- self .assertEqual (asyncio .iscoroutine (iterator .__aiter__ ),
502- asyncio .iscoroutine (mock_iterator .__aiter__ ))
503- self .assertEqual (asyncio .iscoroutine (iterator .__anext__ ),
504- asyncio .iscoroutine (mock_iterator .__anext__ ))
505553
506554 def test_mock_async_for (self ):
507555 async def iterate (iterator ):
@@ -512,19 +560,30 @@ async def iterate(iterator):
512560 return accumulator
513561
514562 expected = ["FOO" , "BAR" , "BAZ" ]
515- with self .subTest ("iterate through default value" ):
516- mock_instance = MagicMock (self .WithAsyncIterator ())
517- self .assertEqual ([], asyncio .run (iterate (mock_instance )))
563+ def test_default (mock_type ):
564+ mock_instance = mock_type (self .WithAsyncIterator ())
565+ self .assertEqual (asyncio .run (iterate (mock_instance )), [])
566+
518567
519- with self . subTest ( "iterate through set return_value" ):
520- mock_instance = MagicMock (self .WithAsyncIterator ())
568+ def test_set_return_value ( mock_type ):
569+ mock_instance = mock_type (self .WithAsyncIterator ())
521570 mock_instance .__aiter__ .return_value = expected [:]
522- self .assertEqual (expected , asyncio .run (iterate (mock_instance )))
571+ self .assertEqual (asyncio .run (iterate (mock_instance )), expected )
523572
524- with self . subTest ( "iterate through set return_value iterator" ):
525- mock_instance = MagicMock (self .WithAsyncIterator ())
573+ def test_set_return_value_iter ( mock_type ):
574+ mock_instance = mock_type (self .WithAsyncIterator ())
526575 mock_instance .__aiter__ .return_value = iter (expected [:])
527- self .assertEqual (expected , asyncio .run (iterate (mock_instance )))
576+ self .assertEqual (asyncio .run (iterate (mock_instance )), expected )
577+
578+ for mock_type in [AsyncMock , MagicMock ]:
579+ with self .subTest (f"default value with { mock_type } " ):
580+ test_default (mock_type )
581+
582+ with self .subTest (f"set return_value with { mock_type } " ):
583+ test_set_return_value (mock_type )
584+
585+ with self .subTest (f"set return_value iterator with { mock_type } " ):
586+ test_set_return_value_iter (mock_type )
528587
529588
530589class AsyncMockAssert (unittest .TestCase ):
0 commit comments