def AsyncContextManager()

in later/unittest/mock.py [0:0]

• 12 lines of code
• 3 McCabe index (conditional complexity)

def AsyncContextManager(return_value: Any = DEFAULT, instance: bool = False) -> Mock:
    """
    This helper sets up a Mock Tree to mock out an AsyncContextManager class.
    return_value: The object returned by the context manager
    instance: Instead of a Class, this mock represents and instance


    mock.assert_awaited() -> The context manager was __aentered__
    mock.assert_not_awaited() -> the context manager was not __aentered__
    """
    mock = Mock()
    manager_instance = AsyncMock()

    if instance:
        mock = manager_instance
    else:
        mock.return_value = manager_instance

    if return_value is not DEFAULT:
        manager_instance.__aenter__.return_value = return_value

    mock.assert_awaited = manager_instance.__aenter__.assert_awaited
    mock.assert_not_awaited = manager_instance.__aenter__.assert_not_awaited
    return mock