Skip to content

Commit

Permalink
fix agents with multiple topics (#116)
Browse files Browse the repository at this point in the history
Co-authored-by: Stevan Milic <stevan.milic@tradecore.com>
  • Loading branch information
stevanmilic and Stevan Milic authored Feb 27, 2021
1 parent 5b68580 commit 70e5516
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 27 deletions.
19 changes: 1 addition & 18 deletions faust/agents/agent.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
"""Agent implementation."""
import asyncio
import sys
import typing
from contextlib import suppress
from contextvars import ContextVar
Expand Down Expand Up @@ -659,23 +658,7 @@ async def _prepare_actor(self, aref: ActorRefT, beacon: NodeT) -> ActorRefT:
else:
# agent yields and is an AsyncIterator so we have to consume it.
coro = self._slurp(aref, aiter(aref))
req_version = (3, 8)
cur_version = sys.version_info
if cur_version >= req_version:
if isinstance(self.channel, TopicT):
name = self.channel.get_topic_name()
else:
name = "channel"
task = asyncio.Task(
self._execute_actor(coro, aref),
loop=self.loop,
name=f"{str(aref)}-{name}",
)
else:
task = asyncio.Task(
self._execute_actor(coro, aref),
loop=self.loop,
)
task = asyncio.Task(self._execute_actor(coro, aref), loop=self.loop)
task._beacon = beacon # type: ignore
aref.actor_task = task
self._actors.add(aref)
Expand Down
28 changes: 19 additions & 9 deletions tests/unit/agents/test_agent.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import asyncio
import sys

import pytest
from mode import SupervisorStrategy, label
Expand Down Expand Up @@ -404,14 +403,7 @@ async def test_prepare_actor__AsyncIterable(self, *, agent):
agent._slurp.assert_called()
coro = agent._slurp()
agent._execute_actor.assert_called_once_with(coro, aref)
if sys.version_info >= (3, 8):
Task.assert_called_once_with(
agent._execute_actor(),
loop=agent.loop,
name=f"{ret}-testid-tests.unit.agents.test_agent.myagent",
)
else:
Task.assert_called_once_with(agent._execute_actor(), loop=agent.loop)
Task.assert_called_once_with(agent._execute_actor(), loop=agent.loop)
task = Task()
assert task._beacon is beacon
assert aref.actor_task is task
Expand All @@ -436,6 +428,24 @@ async def test_prepare_actor__Awaitable(self, *, agent2):
assert aref in agent2._actors
assert ret is aref

@pytest.mark.asyncio
async def test_prepare_actor__Awaitable_with_multiple_topics(self, *, agent2):
aref = agent2(index=0, active_partitions=None)
asyncio.ensure_future(aref.it).cancel() # silence warning
agent2.channel.topics = ["foo", "bar"]
with patch("asyncio.Task") as Task:
agent2._execute_actor = Mock(name="_execute_actor")
beacon = Mock(name="beacon", autospec=Node)
ret = await agent2._prepare_actor(aref, beacon)
coro = aref
agent2._execute_actor.assert_called_once_with(coro, aref)
Task.assert_called_once_with(agent2._execute_actor(), loop=agent2.loop)
task = Task()
assert task._beacon is beacon
assert aref.actor_task is task
assert aref in agent2._actors
assert ret is aref

@pytest.mark.asyncio
async def test_prepare_actor__Awaitable_cannot_have_sinks(self, *, agent2):
aref = agent2(index=0, active_partitions=None)
Expand Down

0 comments on commit 70e5516

Please sign in to comment.