diff --git a/faust/agents/agent.py b/faust/agents/agent.py index 036b0c46d..c78e1b462 100644 --- a/faust/agents/agent.py +++ b/faust/agents/agent.py @@ -4,6 +4,7 @@ import typing from contextlib import suppress from contextvars import ContextVar +from inspect import isasyncgenfunction from time import time from typing import ( Any, @@ -1101,11 +1102,28 @@ def __init__( self.results = {} self.new_value_processed = asyncio.Condition() self.original_channel = cast(ChannelT, original_channel) - self.add_sink(self._on_value_processed) 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) + 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. + # Record the incoming value and wake up any ``put(wait=True)`` caller + # instead, then hand the value on to the agent unchanged. + await self._on_value_processed(value) + return value + async def on_stop(self) -> None: await self._stream.stop() await super().on_stop() diff --git a/tests/unit/agents/test_agent.py b/tests/unit/agents/test_agent.py index 25f033ac6..b94136700 100644 --- a/tests/unit/agents/test_agent.py +++ b/tests/unit/agents/test_agent.py @@ -968,3 +968,26 @@ def dummy_sink(_): async with agent.test_context() as agent_mock: with pytest.raises(SinkCalledException): await agent_mock.put("hello") + + @pytest.mark.skipif( + platform.python_implementation() == "PyPy", reason="Not yet supported on PyPy" + ) + async def test_context__sinkless_agent(self, *, app): + # An agent that never yields cannot use sinks, so test_context() used + # to raise ImproperlyConfigured("Agent must yield to use sinks") at + # startup. It should now work and still let put(wait=True) return. + # See issue #433. + processed = [] + + @app.agent() + async def sinkless(stream): + async for value in stream: + processed.append(value) + + async with sinkless.test_context() as agent: + assert not agent._agent_yields + event = await agent.put("hello") + assert event.value == "hello" + + assert processed == ["hello"] + assert agent.results[0] == "hello"