Fix #433: allow test_context() for sink-less (non-yielding) agents

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()

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"