feat(when): add ContextManager mocking support (#93)

Author: mcous

CONTRIBUTING.md

@@ -64,7 +64,7 @@ poetry run mkdocs serve
 
 The library and documentation will be deployed to PyPI and GitHub Pages, respectively, by CI. To trigger the deploy, cut a new version and push it to GitHub.
 
-Deploy adheres to [semantic versioning][], so care should be taken to bump accurately.
+Decoy adheres to [semantic versioning][], so care should be taken to bump accurately.
 
 ```bash
 # checkout the main branch and pull down latest changes

decoy/__init__.py

@@ -1,21 +1,27 @@
 """Decoy stubbing and spying library."""
-from typing import Any, Callable, Generic, Optional, cast, overload
-
-from . import matchers, errors, warnings
+from typing import Any, Callable, Generic, Optional, Union, cast, overload
+
+from . import errors, matchers, warnings
+from .context_managers import (
+    AsyncContextManager,
+    ContextManager,
+    GeneratorContextManager,
+)
 from .core import DecoyCore, StubCore
-from .types import ClassT, FuncT, ReturnT
+from .types import ClassT, ContextValueT, FuncT, ReturnT
 
 # ensure decoy does not pollute pytest tracebacks
 __tracebackhide__ = True
 
 
 class Decoy:
-    """Decoy test double state container."""
+    """Decoy mock factory and state container."""
 
     def __init__(self) -> None:
-        """Initialize the state container for test doubles and stubs.
+        """Initialize a new mock factory.
 
-        You should initialize a new Decoy instance for every test. See the
+        You should create a new Decoy instance for every test. If you use
+        the `decoy` pytest fixture, this is done automatically. See the
         [setup guide](../#setup) for more details.
         """
         self._core = DecoyCore()
@@ -111,7 +117,8 @@ def when(
                 ignoring unspecified arguments.
 
         Returns:
-            A stub to configure using `then_return`, `then_raise`, or `then_do`.
+            A stub to configure using `then_return`, `then_raise`, `then_do`, or
+            `then_enter_with`.
 
         Example:
             ```python
@@ -137,7 +144,7 @@ def verify(
         times: Optional[int] = None,
         ignore_extra_args: bool = False,
     ) -> None:
-        """Verify a decoy was called using one or more rehearsals.
+        """Verify a mock was called using one or more rehearsals.
 
         See [verification usage guide](../usage/verify/) for more details.
 
@@ -175,11 +182,11 @@ def test_create_something(decoy: Decoy):
         )
 
     def reset(self) -> None:
-        """Reset all decoy state.
+        """Reset all mock state.
 
         This method should be called after every test to ensure spies and stubs
-        don't leak between tests. The Decoy fixture provided by the pytest plugin
-        will do this automatically.
+        don't leak between tests. The `decoy` fixture provided by the pytest plugin
+        will call `reset` automatically.
 
         The `reset` method may also trigger warnings if Decoy detects any questionable
         mock usage. See [decoy.warnings][] for more details.
@@ -228,5 +235,47 @@ def then_do(self, action: Callable[..., ReturnT]) -> None:
         """
         self._core.then_do(action)
 
+    @overload
+    def then_enter_with(
+        self: "Stub[ContextManager[ContextValueT]]",
+        value: ContextValueT,
+    ) -> None:
+        ...
+
+    @overload
+    def then_enter_with(
+        self: "Stub[AsyncContextManager[ContextValueT]]",
+        value: ContextValueT,
+    ) -> None:
+        ...
+
+    @overload
+    def then_enter_with(
+        self: "Stub[GeneratorContextManager[ContextValueT]]",
+        value: ContextValueT,
+    ) -> None:
+        ...
+
+    def then_enter_with(
+        self: Union[
+            "Stub[GeneratorContextManager[ContextValueT]]",
+            "Stub[ContextManager[ContextValueT]]",
+            "Stub[AsyncContextManager[ContextValueT]]",
+        ],
+        value: ContextValueT,
+    ) -> None:
+        """Configure the stub to return a value wrapped in a context manager.
+
+        The wrapping context manager is compatible with both the synchronous and
+        asynchronous context manager interfaces.
+
+        See the [context manager usage guide](../advanced/context-managers/)
+        for more details.
+
+        Arguments:
+            value: A return value to wrap in a ContextManager.
+        """
+        self._core.then_enter_with(value)
+
 
 __all__ = ["Decoy", "Stub", "matchers", "warnings", "errors"]

decoy/call_handler.py

@@ -2,8 +2,9 @@
 from typing import Any
 
 from .call_stack import CallStack
-from .stub_store import StubStore
+from .context_managers import ContextWrapper
 from .spy_calls import SpyCall
+from .stub_store import StubStore
 
 
 class CallHandler:
@@ -25,4 +26,7 @@ def handle(self, call: SpyCall) -> Any:
         if behavior.action:
             return behavior.action(*call.args, **call.kwargs)
 
+        if behavior.context_value:
+            return ContextWrapper(behavior.context_value)
+
         return behavior.return_value

decoy/context_managers.py

@@ -0,0 +1,46 @@
+"""Wrappers around contextlib types and fallbacks."""
+import contextlib
+from typing import Any, AsyncContextManager, ContextManager, Generic, TypeVar
+
+GeneratorContextManager = contextlib._GeneratorContextManager
+
+_EnterT = TypeVar("_EnterT")
+
+
+class ContextWrapper(
+    ContextManager[_EnterT],
+    AsyncContextManager[_EnterT],
+    Generic[_EnterT],
+):
+    """A simple, do-nothing ContextManager that wraps a given value.
+
+    Adapted from `contextlib.nullcontext` to ensure support across
+    all Python versions.
+    """
+
+    def __init__(self, enter_result: _EnterT) -> None:
+        self._enter_result = enter_result
+
+    def __enter__(self) -> _EnterT:
+        """Return the wrapped value."""
+        return self._enter_result
+
+    def __exit__(self, *args: Any, **kwargs: Any) -> Any:
+        """No-op on exit."""
+        pass
+
+    async def __aenter__(self) -> _EnterT:
+        """Return the wrapped value."""
+        return self._enter_result
+
+    async def __aexit__(self, *args: Any, **kwargs: Any) -> Any:
+        """No-op on exit."""
+        pass
+
+
+__all__ = [
+    "AsyncContextManager",
+    "GeneratorContextManager",
+    "ContextManager",
+    "ContextWrapper",
+]

decoy/core.py

@@ -1,14 +1,15 @@
 """Decoy implementation logic."""
 from typing import Any, Callable, Optional
 
-from .spy import SpyConfig, SpyFactory, create_spy as default_create_spy
-from .spy_calls import WhenRehearsal
-from .call_stack import CallStack
-from .stub_store import StubStore, StubBehavior
 from .call_handler import CallHandler
+from .call_stack import CallStack
+from .spy import SpyConfig, SpyFactory
+from .spy import create_spy as default_create_spy
+from .spy_calls import WhenRehearsal
+from .stub_store import StubBehavior, StubStore
+from .types import ContextValueT, ReturnT
 from .verifier import Verifier
 from .warning_checker import WarningChecker
-from .types import ReturnT
 
 # ensure decoy.core does not pollute Pytest tracebacks
 __tracebackhide__ = True
@@ -115,3 +116,10 @@ def then_do(self, action: Callable[..., ReturnT]) -> None:
             rehearsal=self._rehearsal,
             behavior=StubBehavior(action=action),
         )
+
+    def then_enter_with(self, value: ContextValueT) -> None:
+        """Set the stub to return a ContextManager wrapped value."""
+        self._stub_store.add(
+            rehearsal=self._rehearsal,
+            behavior=StubBehavior(context_value=value),
+        )

decoy/pytest_plugin.py

@@ -4,12 +4,14 @@
 fixture without modifying any other pytest behavior. Its usage is optional
 but highly recommended.
 """
-import pytest
 from typing import Iterable
+
+import pytest
+
 from decoy import Decoy
 
 
-@pytest.fixture
+@pytest.fixture()
 def decoy() -> Iterable[Decoy]:
     """Get a [decoy.Decoy][] container and tear it down after the test.
 

decoy/spy.py

@@ -6,7 +6,18 @@
 from inspect import getattr_static, isclass, iscoroutinefunction, isfunction, signature
 from functools import partial
 from warnings import warn
-from typing import get_type_hints, Any, Callable, Dict, NamedTuple, Optional
+from types import TracebackType
+from typing import (
+    cast,
+    get_type_hints,
+    Any,
+    Callable,
+    ContextManager,
+    Dict,
+    NamedTuple,
+    Optional,
+    Type,
+)
 
 from .spy_calls import SpyCall
 from .warnings import IncorrectCallWarning
@@ -37,7 +48,7 @@ def _get_type_hints(obj: Any) -> Dict[str, Any]:
         return {}
 
 
-class BaseSpy:
+class BaseSpy(ContextManager[Any]):
     """Spy object base class.
 
     - Pretends to be another class, if another class is given as a spec
@@ -84,25 +95,35 @@ def __class__(self) -> Any:
 
         return type(self)
 
-    def _call(self, *args: Any, **kwargs: Any) -> Any:
-        spy_id = id(self)
-        spy_name = (
-            self._name
-            if self._name
-            else f"{type(self).__module__}.{type(self).__qualname__}"
-        )
+    def __enter__(self) -> Any:
+        """Allow a spy to be used as a context manager."""
+        enter_spy = self._get_or_create_child_spy("__enter__")
+        return enter_spy()
 
-        if hasattr(self, "__signature__"):
-            try:
-                bound_args = self.__signature__.bind(*args, **kwargs)
-            except TypeError as e:
-                # stacklevel: 3 ensures warning is linked to call location
-                warn(IncorrectCallWarning(e), stacklevel=3)
-            else:
-                args = bound_args.args
-                kwargs = bound_args.kwargs
-
-        return self._handle_call(SpyCall(spy_id, spy_name, args, kwargs))
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Optional[bool]:
+        """Allow a spy to be used as a context manager."""
+        exit_spy = self._get_or_create_child_spy("__exit__")
+        return cast(Optional[bool], exit_spy(exc_type, exc_value, traceback))
+
+    async def __aenter__(self) -> Any:
+        """Allow a spy to be used as an async context manager."""
+        enter_spy = self._get_or_create_child_spy("__aenter__")
+        return await enter_spy()
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Optional[bool]:
+        """Allow a spy to be used as a context manager."""
+        exit_spy = self._get_or_create_child_spy("__aexit__")
+        return cast(Optional[bool], await exit_spy(exc_type, exc_value, traceback))
 
     def __repr__(self) -> str:
         """Get a helpful string representation of the spy."""
@@ -118,14 +139,15 @@ def __repr__(self) -> str:
             return "<Decoy mock>"
 
     def __getattr__(self, name: str) -> Any:
-        """Get a property of the spy.
-
-        Lazily constructs child spies, basing them on type hints if available.
-        """
+        """Get a property of the spy, always returning a child spy."""
         # do not attempt to mock magic methods
         if name.startswith("__") and name.endswith("__"):
             return super().__getattribute__(name)
 
+        return self._get_or_create_child_spy(name)
+
+    def _get_or_create_child_spy(self, name: str) -> Any:
+        """Lazily construct a child spy, basing it on type hints if available."""
         # return previously constructed (and cached) child spies
         if name in self._spy_children:
             return self._spy_children[name]
@@ -167,6 +189,26 @@ def __getattr__(self, name: str) -> Any:
 
         return spy
 
+    def _call(self, *args: Any, **kwargs: Any) -> Any:
+        spy_id = id(self)
+        spy_name = (
+            self._name
+            if self._name
+            else f"{type(self).__module__}.{type(self).__qualname__}"
+        )
+
+        if hasattr(self, "__signature__"):
+            try:
+                bound_args = self.__signature__.bind(*args, **kwargs)
+            except TypeError as e:
+                # stacklevel: 3 ensures warning is linked to call location
+                warn(IncorrectCallWarning(e), stacklevel=3)
+            else:
+                args = bound_args.args
+                kwargs = bound_args.kwargs
+
+        return self._handle_call(SpyCall(spy_id, spy_name, args, kwargs))
+
 
 class Spy(BaseSpy):
     """An object that records all calls made to itself and its children."""

decoy/stub_store.py

@@ -8,6 +8,7 @@ class StubBehavior(NamedTuple):
     """A recorded stub behavior."""
 
     return_value: Optional[Any] = None
+    context_value: Optional[Any] = None
     error: Optional[Exception] = None
     action: Optional[Callable[..., Any]] = None
     once: bool = False

decoy/types.py

@@ -9,3 +9,6 @@
 
 ReturnT = TypeVar("ReturnT")
 """The return type of a given call."""
+
+ContextValueT = TypeVar("ContextValueT")
+"""A context manager value returned by a stub."""