Fix test_context() recording input for callable-object agents

CHANGELOG.md

@@ -53,6 +53,14 @@ resumes the Keep a Changelog format.
 - `faust[aerospike]` installed nothing: `requirements/extras/aerospike.txt`
   shipped without the matching `BUNDLES` entry in `setup.py`, despite being
   advertised in the README. A new test guards both directions of that mapping.
+- `agent.test_context()` records what a yielding agent yields even when its
+  function is not itself an async-generator function — a callable object with
+  an async-generator `__call__`, or a function returning an async generator.
+  Previously the wrapper classified the agent with
+  `inspect.isasyncgenfunction()`, which is `False` for both shapes, so they
+  were treated as never yielding and `agent.results` was filled with the values
+  sent *in* rather than the values yielded. Nothing raised, so tests kept
+  passing while asserting against their own input.
 - Every Kafka rebalance failed when `opentracing` was not installed. The no-op
   stand-in Faust falls back to gave its spans no tracer, but
   `traced_from_parent_span` starts a child span from `parent.tracer`, so

docs/userguide/testing.rst

@@ -389,6 +389,12 @@ yields                the value the agent **yielded** (its output)
 does not yield        the value that was **sent in** (its input)
 ===================== =============================================
 
+Which row applies is decided from the running agent, not from how its function
+is written.  An agent implemented as a callable object with an
+async-generator ``__call__``, or as a plain function returning an async
+generator, yields just as much as an ``async def`` that yields directly, and
+its ``results`` hold its output.
+
 There is no output to capture for a sink-less agent, so faust records the
 incoming value instead.  That still makes ``results`` useful for confirming
 which values reached the agent -- as ``test_results_records_input_values``

faust/agents/agent.py

@@ -4,7 +4,6 @@
 import typing
 from contextlib import suppress
 from contextvars import ContextVar
-from inspect import isasyncgenfunction
 from time import time
 from typing import (
     Any,
@@ -1117,6 +1116,10 @@ def shortlabel(self) -> str:
 class AgentTestWrapper(Agent, AgentTestWrapperT):  # pragma: no cover
     _stream: StreamT
 
+    # None until the first actor is built, since whether the agent yields is
+    # only knowable from the actor itself.  See ``_prepare_actor``.
+    _agent_yields: Optional[bool] = None
+
     def __init__(
         self, *args: Any, original_channel: Optional[ChannelT] = None, **kwargs: Any
     ) -> None:
@@ -1125,19 +1128,26 @@ def __init__(
         self.new_value_processed = asyncio.Condition()
         self.original_channel = cast(ChannelT, original_channel)
         self._stream = self.channel.stream()
-        # Agents that never yield cannot use sinks -- ``_prepare_actor``
-        # raises ``ImproperlyConfigured('Agent must yield to use sinks')``
-        # for them.  So only attach the results sink when the wrapped agent
-        # actually yields; for sink-less agents observe processed values with
-        # a stream processor instead, so ``test_context`` works either way.
-        # See issue #433.
-        self._agent_yields = isasyncgenfunction(self.fun)
+        self.sent_offset = 0
+        self.processed_offset = 0
+
+    async def _prepare_actor(self, aref: ActorRefT, beacon: NodeT) -> ActorRefT:
+        self._install_value_recorder(aref)
+        return await super()._prepare_actor(aref, beacon)
+
+    def _install_value_recorder(self, aref: ActorRefT) -> None:
+        # Install once per wrapper, since all actors share one stream chain and
+        # ``Stream.add_processor`` is not idempotent, and before
+        # ``super()._prepare_actor`` starts the task, since deriving a stream
+        # via ``group_by``/``through`` moves processors off ``self._stream``.
+        if self._agent_yields is not None:
+            return
+        # Mirrors ``Agent._prepare_actor``, which rejects sinks when not yielding.
+        self._agent_yields = not isinstance(aref, Awaitable)
         if self._agent_yields:
             self.add_sink(self._on_value_processed)
         else:
             self._stream.add_processor(self._on_value_processed_processor)
-        self.sent_offset = 0
-        self.processed_offset = 0
 
     async def _on_value_processed_processor(self, value: Any) -> Any:
         # Sink-less agents don't yield, so we can't observe their output.

tests/unit/agents/test_agent.py

@@ -1,4 +1,5 @@
 import asyncio
+import functools
 from unittest.mock import ANY, call, patch
 
 import pytest
@@ -1000,9 +1001,74 @@ async def sinkless(stream):
                 processed.append(value)
 
         async with sinkless.test_context() as agent:
-            assert not agent._agent_yields
+            assert agent._agent_yields is False
             event = await agent.put("hello")
             assert event.value == "hello"
 
         assert processed == ["hello"]
         assert agent.results[0] == "hello"
+
+    async def test_context__callable_object_agent(self, *, app):
+        # ``isasyncgenfunction`` returns False for a class instance whose
+        # ``__call__`` is an async generator, but the agent does yield.
+        # ``results`` must hold the yielded value, not the incoming one.
+        class CustomAgent:
+            async def __call__(self, stream):
+                async for value in stream:
+                    yield value.upper()
+
+        my_agent = app.agent(name="callable_object")(CustomAgent())
+
+        async with my_agent.test_context() as agent:
+            await agent.put("hello")
+            assert agent.results[0] == "HELLO"
+            assert agent._agent_yields is True
+
+    async def test_context__function_returning_async_generator(self, *, app):
+        # A plain ``def`` that returns an async generator object also yields,
+        # and is likewise invisible to ``isasyncgenfunction``.
+        async def transform(stream):
+            async for value in stream:
+                yield value.upper()
+
+        def returns_async_generator(stream):
+            return transform(stream)
+
+        my_agent = app.agent(name="returns_agen")(returns_async_generator)
+
+        async with my_agent.test_context() as agent:
+            await agent.put("hello")
+            assert agent.results[0] == "HELLO"
+            assert agent._agent_yields is True
+
+    async def test_context__partial_over_callable_object(self, *, app):
+        # ``isasyncgenfunction`` sees through a partial over an async generator
+        # function, but not over a callable object.  ``name`` is explicit
+        # because every ``functools.partial`` shares the same derived name.
+        class CustomAgent:
+            async def __call__(self, stream):
+                async for value in stream:
+                    yield value.upper()
+
+        my_agent = app.agent(name="partial_over_callable")(
+            functools.partial(CustomAgent())
+        )
+
+        async with my_agent.test_context() as agent:
+            await agent.put("hello")
+            assert agent.results[0] == "HELLO"
+            assert agent._agent_yields is True
+
+    async def test_context__records_each_value_once_with_concurrency(self, *, app):
+        # Every actor of one test wrapper resolves against the same stream
+        # chain, and ``Stream.add_processor`` is not idempotent, so recording
+        # must be installed once per wrapper rather than once per actor.
+        @app.agent(name="concurrent_sinkless", concurrency=3)
+        async def sinkless(stream):
+            async for value in stream:
+                pass
+
+        async with sinkless.test_context() as agent:
+            await agent.put("hello")
+            assert len(agent.results) == 1
+            assert agent._agent_yields is False