Validate contextmanager return types against the context manager

tests/lib_testslide.py

@@ -376,6 +376,32 @@ def passes_for_valid_forward_reference(self):
             Foo.get_maybe_foo, Foo(), self.caller_frame_info
         )
 
+    @context.example
+    def passes_for_context_manager_template(self):
+        """
+        contextlib.contextmanager copies __annotations__ from the generator
+        function it wraps, so the recorded return type describes what the
+        generator yields rather than the context manager that calling it
+        actually returns.
+        https://github.com/facebook/TestSlide/issues/193
+        """
+        self.callable_template = sample_module.test_function_returns_context_manager
+        self.assert_passes(StrictMock(template=sample_module.ContextManagerTarget))
+
+    @context.example
+    def passes_for_async_context_manager_template(self):
+        """Same as above, for contextlib.asynccontextmanager."""
+        self.callable_template = (
+            sample_module.test_function_returns_async_context_manager
+        )
+        self.assert_passes(StrictMock(template=sample_module.ContextManagerTarget))
+
+    @context.example
+    def fails_for_context_manager_template_given_a_non_context_manager(self):
+        """The check is redirected, not disabled."""
+        self.callable_template = sample_module.test_function_returns_context_manager
+        self.assert_fails(42)
+
     @context.example
     def fails_for_valid_forward_reference_but_bad_type_passed(self):
         with self.assertRaisesRegex(

tests/sample_module.py

@@ -3,9 +3,12 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from collections.abc import Awaitable, Coroutine
+from collections.abc import AsyncGenerator, Awaitable, Coroutine, Generator
+from contextlib import asynccontextmanager, contextmanager
 from typing import Any, Union
 
+from typing_extensions import Self
+
 attribute = "value"
 typedattr: str = "bruh"
 
@@ -175,3 +178,35 @@ def test_union(arg: UnionArgType) -> None:
 
 def test_tuple(arg: TupleArgType) -> None:
     pass
+
+
+class ContextManagerTarget:
+    "This class is used by some unit tests only"
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, exc_type: object, exc: object, tb: object) -> None:
+        return None
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
+        return None
+
+
+@contextmanager
+def test_function_returns_context_manager() -> Generator[
+    "ContextManagerTarget", None, None
+]:
+    "This function is used by some unit tests only"
+    yield ContextManagerTarget()
+
+
+@asynccontextmanager
+async def test_function_returns_async_context_manager() -> AsyncGenerator[
+    "ContextManagerTarget", None
+]:
+    "This function is used by some unit tests only"
+    yield ContextManagerTarget()

testslide/core/lib.py

@@ -5,6 +5,7 @@
 
 # pyre-unsafe
 import collections.abc as abc
+import contextlib
 import functools
 import inspect
 import os
@@ -341,6 +342,33 @@ def _is_wrapped_for_signature_and_type_validation(value: Callable) -> bool:
     return getattr(value, "__is_testslide_type_validation_wrapping", False)
 
 
+def _get_contextmanager_return_type(template: Any) -> Any:
+    """Return the context manager type produced by a contextlib-decorated function.
+
+    ``contextlib.contextmanager`` and ``contextlib.asynccontextmanager`` wrap the
+    decorated generator function with ``functools.wraps``, which copies
+    ``__annotations__`` over verbatim. The recorded return type therefore
+    describes what the *generator* yields (eg ``AsyncGenerator[Foo, None]``),
+    while calling the decorated function actually returns a context manager.
+    Validating a mocked return value against the generator type rejects every
+    legitimate value.
+
+    Returns ``None`` when ``template`` is not such a decorated function.
+    """
+    wrapped = getattr(template, "__wrapped__", None)
+    if wrapped is None:
+        return None
+    # The decorator turns a generator function into one that is no longer a
+    # generator function, which is what distinguishes it from @wraps in general.
+    if inspect.isasyncgenfunction(wrapped) and not inspect.isasyncgenfunction(template):
+        return contextlib.AbstractAsyncContextManager
+    if inspect.isgeneratorfunction(wrapped) and not inspect.isgeneratorfunction(
+        template
+    ):
+        return contextlib.AbstractContextManager
+    return None
+
+
 def _validate_return_type(
     template: Mock | Callable,
     value: Any,
@@ -352,6 +380,9 @@ def _validate_return_type(
     except TypeError:
         return
     expected_type = argspec.annotations.get("return")
+    contextmanager_type = _get_contextmanager_return_type(template)
+    if contextmanager_type is not None:
+        expected_type = contextmanager_type
     if expected_type:
         if unwrap_template_awaitable:
             type_origin = get_origin(expected_type)