From 6e1024cdcf48bd96337e5e7db8cea0c8095aae6d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Burczy=C5=84ski?= Date: Wed, 19 Aug 2026 10:16:08 +0200 Subject: [PATCH] Fix test_context() recording input for callable-object agents Fixes #784. --- CHANGELOG.md | 8 ++++ docs/userguide/testing.rst | 6 +++ faust/agents/agent.py | 30 ++++++++++----- tests/unit/agents/test_agent.py | 68 ++++++++++++++++++++++++++++++++- 4 files changed, 101 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1f9524dce..34abadf6e 100644 --- a/CHANGELOG.md +++ b/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. ### Changed - The `examples/fastapi/` directory is now `examples/fastapi_project/`. The old diff --git a/docs/userguide/testing.rst b/docs/userguide/testing.rst index 5db3060dd..34e812965 100644 --- a/docs/userguide/testing.rst +++ b/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`` diff --git a/faust/agents/agent.py b/faust/agents/agent.py index 76a0a0631..cf5554439 100644 --- a/faust/agents/agent.py +++ b/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. diff --git a/tests/unit/agents/test_agent.py b/tests/unit/agents/test_agent.py index 96f5c316e..2c36b150f 100644 --- a/tests/unit/agents/test_agent.py +++ b/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