예제 #1
0
    async def test_slurp(self, *, agent, app):
        aref = agent(index=None, active_partitions=None)
        stream = aref.stream.get_active_stream()
        agent._delegate_to_sinks = AsyncMock(name='_delegate_to_sinks')
        agent._reply = AsyncMock(name='_reply')

        def on_delegate(value):
            raise StopAsyncIteration()

        word = Word('word')
        word_req = ReqRepRequest(word, 'reply_to', 'correlation_id')
        message1 = Mock(name='message1', autospec=Message)
        message2 = Mock(name='message2', autospec=Message)
        event1 = Event(app, None, word_req, message1)
        event2 = Event(app, 'key', 'bar', message2)
        values = [
            (event1, word),
            (event2, 'bar'),
        ]

        class AIT:
            async def __aiter__(self):
                for event, value in values:
                    stream.current_event = event
                    yield value

        it = aiter(AIT())
        await agent._slurp(aref, it)

        agent._reply.assert_called_once_with(None, word, word_req)
        agent._delegate_to_sinks.coro.assert_has_calls([
            call(word),
            call('bar'),
        ])
예제 #2
0
파일: test_agent.py 프로젝트: taybin/faust
    async def test_slurp(self, *, agent, app):
        aref = agent(index=None, active_partitions=None)
        stream = aref.stream.get_active_stream()
        agent._delegate_to_sinks = AsyncMock(name="_delegate_to_sinks")
        agent._reply = AsyncMock(name="_reply")

        def on_delegate(value):
            raise StopAsyncIteration()

        word = Word("word")
        word_req = ReqRepRequest(word, "reply_to", "correlation_id")
        message1 = Mock(name="message1", autospec=Message)
        message2 = Mock(name="message2", autospec=Message)
        event1 = Event(app, None, word_req, {}, message1)
        event2 = Event(app, "key", "bar", {}, message2)
        values = [
            (event1, word),
            (event2, "bar"),
        ]

        class AIT:
            async def __aiter__(self):
                for event, value in values:
                    stream.current_event = event
                    yield value

        it = aiter(AIT())
        await agent._slurp(aref, it)

        agent._reply.assert_called_once_with(None, word, word_req.reply_to,
                                             word_req.correlation_id)
        agent._delegate_to_sinks.coro.assert_has_calls([
            call(word),
            call("bar"),
        ])
예제 #3
0
 async def test_reply(self, *, agent):
     agent.app = Mock(
         name='app',
         autospec=App,
         send=AsyncMock(),
     )
     req = ReqRepRequest('value', 'reply_to', 'correlation_id')
     await agent._reply('key', 'reply', req)
     agent.app.send.assert_called_once_with(
         req.reply_to,
         key=None,
         value=ReqRepResponse(
             key='key',
             value='reply',
             correlation_id=req.correlation_id,
         ),
     )
예제 #4
0
파일: test_agent.py 프로젝트: taybin/faust
 async def test_reply(self, *, agent):
     agent.app = Mock(
         name="app",
         autospec=App,
         send=AsyncMock(),
     )
     req = ReqRepRequest("value", "reply_to", "correlation_id")
     await agent._reply("key", "reply", req.reply_to, req.correlation_id)
     agent.app.send.assert_called_once_with(
         req.reply_to,
         key=None,
         value=ReqRepResponse(
             key="key",
             value="reply",
             correlation_id=req.correlation_id,
         ),
     )
예제 #5
0
class test_Agent:
    @pytest.fixture
    def agent(self, *, app):
        @app.agent()
        async def myagent(stream):
            async for value in stream:
                yield value

        return myagent

    @pytest.fixture
    def isolated_agent(self, *, app):
        @app.agent(isolated_partitions=True)
        async def isoagent(stream):
            async for value in stream:
                yield value

        return isoagent

    @pytest.fixture
    def foo_topic(self, *, app):
        return app.topic('foo')

    @pytest.fixture
    def agent2(self, *, app, foo_topic):
        @app.agent(foo_topic)
        async def other_agent(stream):
            async for value in stream:
                value

        return other_agent

    def test_init_key_type_and_channel(self, *, app):
        with pytest.raises(AssertionError):

            @app.agent(app.topic('foo'), key_type=bytes)
            async def foo():
                ...

    def test_init_value_type_and_channel(self, *, app):
        with pytest.raises(AssertionError):

            @app.agent(app.topic('foo'), value_type=bytes)
            async def foo():
                ...

    def test_isolated_partitions_cannot_have_concurrency(self, *, app):
        with pytest.raises(ImproperlyConfigured):

            @app.agent(isolated_partitions=True, concurrency=100)
            async def foo():
                ...

    def test_cancel(self, *, agent):
        actor1 = Mock(name='actor1')
        actor2 = Mock(name='actor2')
        agent._actors = [actor1, actor2]
        agent.cancel()
        actor1.cancel.assert_called_once_with()
        actor2.cancel.assert_called_once_with()

    @pytest.mark.asyncio
    async def test_on_partitions_revoked(self, *, agent):
        revoked = {TP('foo', 0)}
        agent.on_shared_partitions_revoked = AsyncMock(name='ospr')
        await agent.on_partitions_revoked(revoked)
        agent.on_shared_partitions_revoked.assert_called_once_with(revoked)

    @pytest.mark.asyncio
    async def test_on_partitions_revoked__isolated(self, *, isolated_agent):
        revoked = {TP('foo', 0)}
        i = isolated_agent.on_isolated_partitions_revoked = AsyncMock(name='i')
        await isolated_agent.on_partitions_revoked(revoked)
        i.assert_called_once_with(revoked)

    @pytest.mark.asyncio
    async def test_on_partitions_assigned(self, *, agent):
        assigned = {TP('foo', 0)}
        agent.on_shared_partitions_assigned = AsyncMock(name='ospr')
        await agent.on_partitions_assigned(assigned)
        agent.on_shared_partitions_assigned.assert_called_once_with(assigned)

    @pytest.mark.asyncio
    async def test_on_partitions_assigned__isolated(self, *, isolated_agent):
        assigned = {TP('foo', 0)}
        i = isolated_agent.on_isolated_partitions_assigned = AsyncMock()
        await isolated_agent.on_partitions_assigned(assigned)
        i.assert_called_once_with(assigned)

    @pytest.mark.asyncio
    async def test_on_isolated_partitions_revoked(self, *, agent):
        tp = TP('foo', 0)
        aref = Mock(
            name='aref',
            autospec=Actor,
            on_isolated_partition_revoked=AsyncMock(),
        )
        agent._actor_by_partition = {tp: aref}

        await agent.on_isolated_partitions_revoked({tp})
        aref.on_isolated_partition_revoked.assert_called_once_with(tp)
        assert not agent._actor_by_partition
        await agent.on_isolated_partitions_revoked({tp})

    @pytest.mark.asyncio
    async def test_on_isolated_partitions_assigned(self, *, agent):
        agent._assign_isolated_partition = AsyncMock(name='aip')
        await agent.on_isolated_partitions_assigned({TP('foo', 0)})
        agent._assign_isolated_partition.assert_called_once_with(TP('foo', 0))

    @pytest.mark.asyncio
    async def test_assign_isolated_partition(self, *, agent):
        agent._on_first_isolated_partition_assigned = Mock(name='ofipa')
        agent._maybe_start_isolated = AsyncMock(name='maybe_start_isolated')
        agent._first_assignment_done = True

        tp = TP('foo', 606)
        await agent._assign_isolated_partition(tp)

        agent._on_first_isolated_partition_assigned.assert_not_called()
        agent._maybe_start_isolated.assert_called_once_with(tp)

        agent._first_assignment_done = False
        agent._actor_by_partition = set()
        await agent._assign_isolated_partition(tp)
        agent._on_first_isolated_partition_assigned.assert_called_once_with(tp)

    def test_on_first_isolated_partition_assigned(self, *, agent):
        aref = Mock(name='actor', autospec=Actor)
        agent._actors = [aref]
        agent._pending_active_partitions = set()
        tp = TP('foo', 303)
        agent._on_first_isolated_partition_assigned(tp)
        assert agent._actor_by_partition[tp] is aref
        assert agent._pending_active_partitions == {tp}
        agent._pending_active_partitions = None
        agent._on_first_isolated_partition_assigned(tp)

    @pytest.mark.asyncio
    async def test_maybe_start_isolated(self, *, isolated_agent):
        aref = Mock(
            name='actor',
            autospec=Actor,
            on_isolated_partition_assigned=AsyncMock(),
        )
        isolated_agent._start_isolated = AsyncMock(
            name='_start_isolated',
            return_value=aref,
        )
        tp = TP('foo', 303)
        await isolated_agent._maybe_start_isolated(tp)

        isolated_agent._start_isolated.assert_called_once_with(tp)
        assert isolated_agent._actor_by_partition[tp] is aref
        aref.on_isolated_partition_assigned.assert_called_once_with(tp)

    @pytest.mark.asyncio
    async def test_start_isolated(self, *, agent):
        service = agent._service = Mock(
            name='service',
            autospec=AgentService,
            _start_for_partitions=AsyncMock(),
        )
        ret = await agent._start_isolated(TP('foo', 0))
        service._start_for_partitions.assert_called_once_with({TP('foo', 0)})
        assert ret is service._start_for_partitions.coro()

    @pytest.mark.asyncio
    async def test_on_shared_partitions_revoked(self, *, agent):
        await agent.on_shared_partitions_revoked(set())

    @pytest.mark.asyncio
    async def test_on_shared_partitions_assigned(self, *, agent):
        await agent.on_shared_partitions_assigned(set())

    def test_info(self, *, agent):
        assert agent.info() == {
            'app': agent.app,
            'fun': agent.fun,
            'name': agent.name,
            'channel': agent.channel,
            'concurrency': agent.concurrency,
            'help': agent.help,
            'sinks': agent._sinks,
            'on_error': agent._on_error,
            'supervisor_strategy': agent.supervisor_strategy,
            'isolated_partitions': agent.isolated_partitions,
        }

    def test_clone(self, *, agent):
        assert agent.clone(isolated_partitions=True).isolated_partitions

    def test_stream__active_partitions(self, *, agent):
        assert agent.stream(active_partitions={TP('foo', 0)})

    @pytest.mark.parametrize('input,expected', [
        (ReqRepRequest('value', 'reply_to', 'correlation_id'), 'value'),
        ('value', 'value'),
    ])
    def test_maybe_unwrap_reply_request(self, input, expected, *, agent):
        assert agent._maybe_unwrap_reply_request(input) == expected

    @pytest.mark.asyncio
    async def test_start_task(self, *, agent):
        agent._prepare_actor = AsyncMock(name='_prepare_actor')
        ret = await agent._start_task(index=0)
        agent._prepare_actor.assert_called_once_with(ANY, None)
        assert ret is agent._prepare_actor.coro()

    @pytest.mark.asyncio
    async def test_prepare_actor__AsyncIterable(self, *, agent):
        aref = agent(index=0, active_partitions=None)
        with patch('asyncio.Task') as Task:
            agent._slurp = Mock(name='_slurp')
            agent._execute_task = Mock(name='_execute_task')
            beacon = Mock(name='beacon', autospec=Node)
            ret = await agent._prepare_actor(aref, beacon)
            agent._slurp.assert_called()
            coro = agent._slurp()
            agent._execute_task.assert_called_once_with(coro, aref)
            Task.assert_called_once_with(agent._execute_task(),
                                         loop=agent.loop)
            task = Task()
            assert task._beacon is beacon
            assert aref.actor_task is task
            assert aref in agent._actors
            assert ret is aref

    @pytest.mark.asyncio
    async def test_prepare_actor__Awaitable(self, *, agent2):
        aref = agent2(index=0, active_partitions=None)
        asyncio.ensure_future(aref.it).cancel()  # silence warning
        return
        with patch('asyncio.Task') as Task:
            agent2._execute_task = Mock(name='_execute_task')
            beacon = Mock(name='beacon', autospec=Node)
            ret = await agent2._prepare_actor(aref, beacon)
            coro = aref
            agent2._execute_task.assert_called_once_with(coro, aref)
            Task.assert_called_once_with(agent2._execute_task(),
                                         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)
        asyncio.ensure_future(aref.it).cancel()  # silence warning
        agent2._sinks = [agent2]
        with pytest.raises(ImproperlyConfigured):
            await agent2._prepare_actor(
                aref,
                Mock(name='beacon', autospec=Node),
            )

    @pytest.mark.asyncio
    async def test_execute_task(self, *, agent):
        coro = done_future()
        await agent._execute_task(coro, Mock(name='aref', autospec=Actor))

    @pytest.mark.asyncio
    async def test_execute_task__cancelled_stopped(self, *, agent):
        coro = FutureMock()
        coro.side_effect = asyncio.CancelledError()
        await agent.stop()
        with pytest.raises(asyncio.CancelledError):
            await agent._execute_task(coro, Mock(name='aref', autospec=Actor))
        coro.assert_awaited()

    @pytest.mark.asyncio
    async def test_execute_task__cancelled_running(self, *, agent):
        coro = FutureMock()
        coro.side_effect = asyncio.CancelledError()
        await agent._execute_task(coro, Mock(name='aref', autospec=Actor))
        coro.assert_awaited()

    @pytest.mark.asyncio
    async def test_execute_task__raising(self, *, agent):
        agent._on_error = AsyncMock(name='on_error')
        agent.log = Mock(name='log', autospec=CompositeLogger)
        aref = Mock(
            name='aref',
            autospec=Actor,
            crash=AsyncMock(),
        )
        agent._service = Mock(name='_service', autospec=AgentService)
        coro = FutureMock()
        exc = coro.side_effect = KeyError('bar')
        with pytest.raises(KeyError):
            await agent._execute_task(coro, aref)
        coro.assert_awaited()

        aref.crash.assert_called_once_with(exc)
        agent._service.supervisor.wakeup.assert_called_once_with()
        agent._on_error.assert_called_once_with(agent, exc)

        agent._on_error = None
        with pytest.raises(KeyError):
            await agent._execute_task(coro, aref)

    @pytest.mark.asyncio
    async def test_slurp(self, *, agent, app):
        aref = agent(index=None, active_partitions=None)
        stream = aref.stream.get_active_stream()
        agent._delegate_to_sinks = AsyncMock(name='_delegate_to_sinks')
        agent._reply = AsyncMock(name='_reply')

        def on_delegate(value):
            raise StopAsyncIteration()

        word = Word('word')
        word_req = ReqRepRequest(word, 'reply_to', 'correlation_id')
        message1 = Mock(name='message1', autospec=Message)
        message2 = Mock(name='message2', autospec=Message)
        event1 = Event(app, None, word_req, message1)
        event2 = Event(app, 'key', 'bar', message2)
        values = [
            (event1, word),
            (event2, 'bar'),
        ]

        class AIT:
            async def __aiter__(self):
                for event, value in values:
                    stream.current_event = event
                    yield value

        it = aiter(AIT())
        await agent._slurp(aref, it)

        agent._reply.assert_called_once_with(None, word, word_req)
        agent._delegate_to_sinks.coro.assert_has_calls([
            call(word),
            call('bar'),
        ])

    @pytest.mark.asyncio
    async def test_delegate_to_sinks(self, *, agent, agent2, foo_topic):
        agent2.send = AsyncMock(name='agent2.send')
        foo_topic.send = AsyncMock(name='foo_topic.send')
        sink_callback = Mock(name='sink_callback')
        sink_callback2_mock = Mock(name='sink_callback2_mock')

        async def sink_callback2(value):
            return sink_callback2_mock(value)

        agent._sinks = [
            agent2,
            foo_topic,
            sink_callback,
            sink_callback2,
        ]

        value = Mock(name='value')
        await agent._delegate_to_sinks(value)

        agent2.send.assert_called_once_with(value=value)
        foo_topic.send.assert_called_once_with(value=value)
        sink_callback.assert_called_once_with(value)
        sink_callback2_mock.assert_called_once_with(value)

    @pytest.mark.asyncio
    async def test_reply(self, *, agent):
        agent.app = Mock(
            name='app',
            autospec=App,
            send=AsyncMock(),
        )
        req = ReqRepRequest('value', 'reply_to', 'correlation_id')
        await agent._reply('key', 'reply', req)
        agent.app.send.assert_called_once_with(
            req.reply_to,
            key=None,
            value=ReqRepResponse(
                key='key',
                value='reply',
                correlation_id=req.correlation_id,
            ),
        )

    @pytest.mark.asyncio
    async def test_cast(self, *, agent):
        agent.send = AsyncMock(name='send')
        await agent.cast('value', key='key', partition=303)
        agent.send.assert_called_once_with(key='key',
                                           value='value',
                                           partition=303)

    @pytest.mark.asyncio
    async def test_ask(self, *, agent):
        agent.app = Mock(
            name='app',
            autospec=App,
            maybe_start_client=AsyncMock(),
            _reply_consumer=Mock(
                autospec=ReplyConsumer,
                add=AsyncMock(),
            ),
        )
        pp = done_future()
        agent.ask_nowait = Mock(name='ask_nowait')
        agent.ask_nowait.return_value = done_future(pp)
        pp.correlation_id = 'foo'

        await agent.ask(
            value='val',
            key='key',
            partition=303,
            correlation_id='correlation_id',
        )
        agent.ask_nowait.assert_called_once_with(
            'val',
            key='key',
            partition=303,
            reply_to=agent.app.conf.reply_to,
            correlation_id='correlation_id',
            force=True,
        )
        agent.app._reply_consumer.add.assert_called_once_with(
            pp.correlation_id, pp)

    @pytest.mark.asyncio
    async def test_ask_nowait(self, *, agent):
        agent._create_req = Mock(name='_create_req')
        agent.channel.send = AsyncMock(name='channel.send')
        res = await agent.ask_nowait(
            value='value',
            key='key',
            partition=303,
            reply_to='reply_to',
            correlation_id='correlation_id',
            force=True,
        )

        agent._create_req.assert_called_once_with('key', 'value', 'reply_to',
                                                  'correlation_id')
        agent.channel.send.assert_called_once_with(key='key',
                                                   value=agent._create_req(),
                                                   partition=303,
                                                   force=True)

        assert res.reply_to == agent._create_req().reply_to
        assert res.correlation_id == agent._create_req().correlation_id

    def test_create_req(self, *, agent):
        agent._get_strtopic = Mock(name='_get_strtopic')
        with patch('faust.agents.agent.uuid4') as uuid4:
            uuid4.return_value = 'vvv'
            reqrep = agent._create_req(key=b'key',
                                       value=b'value',
                                       reply_to='reply_to')

            agent._get_strtopic.assert_called_once_with('reply_to')

            assert reqrep.value == b'value'
            assert reqrep.reply_to == agent._get_strtopic()
            assert reqrep.correlation_id == 'vvv'

    @pytest.mark.asyncio
    async def test_send(self, *, agent):
        agent.channel = Mock(
            name='channel',
            autospec=Channel,
            send=AsyncMock(),
        )
        agent._create_req = Mock(name='_create_req')
        callback = Mock(name='callback')

        ret = await agent.send(
            key=b'key',
            value=b'value',
            partition=303,
            key_serializer='raw',
            value_serializer='raw',
            callback=callback,
            reply_to='reply_to',
            correlation_id='correlation_id',
            force=True,
        )

        agent._create_req.assert_called_once_with(
            b'key',
            b'value',
            'reply_to',
            'correlation_id',
        )

        agent.channel.send.assert_called_once_with(
            key=b'key',
            value=agent._create_req(),
            partition=303,
            key_serializer='raw',
            value_serializer='raw',
            force=True,
        )

        assert ret is agent.channel.send.coro()

    @pytest.mark.asyncio
    async def test_send__without_reply_to(self, *, agent):
        agent.channel = Mock(
            name='channel',
            autospec=Channel,
            send=AsyncMock(),
        )
        agent._create_req = Mock(name='_create_req')
        callback = Mock(name='callback')

        ret = await agent.send(
            key=b'key',
            value=b'value',
            partition=303,
            key_serializer='raw',
            value_serializer='raw',
            callback=callback,
            reply_to=None,
            correlation_id='correlation_id',
            force=True,
        )

        agent._create_req.assert_not_called()

        agent.channel.send.assert_called_once_with(
            key=b'key',
            value=b'value',
            partition=303,
            key_serializer='raw',
            value_serializer='raw',
            force=True,
        )

        assert ret is agent.channel.send.coro()

    def test_get_strtopic__agent(self, *, agent, agent2):
        assert agent._get_strtopic(agent2) == agent2.channel.get_topic_name()

    def test_get_strtopic__topic(self, *, agent, foo_topic):
        assert agent._get_strtopic(foo_topic) == foo_topic.get_topic_name()

    def test_get_strtopic__str(self, *, agent):
        assert agent._get_strtopic('bar') == 'bar'

    def test_get_strtopic__channel_raises(self, *, agent, app):
        with pytest.raises(ValueError):
            agent._get_strtopic(app.channel())

    def test_get_topic_names(self, *, agent, app):
        agent.channel = app.topic('foo')
        assert agent.get_topic_names() == ('foo', )

    def test_get_topic_names__channel(self, *, agent, app):
        agent.channel = app.channel()
        assert agent.get_topic_names() == []

    def test_repr(self, *, agent):
        assert repr(agent)

    def test_channel(self, *, agent):
        agent._prepare_channel = Mock(name='_prepare_channel')
        agent._channel = None
        channel = agent.channel
        agent._prepare_channel.assert_called_once_with(
            agent._channel_arg,
            key_type=agent._key_type,
            value_type=agent._value_type,
            **agent._channel_kwargs)
        assert channel is agent._prepare_channel.return_value
        assert agent._channel is channel

    def test_prepare_channel__not_channel(self, *, agent):
        with pytest.raises(TypeError):
            agent._prepare_channel(object())

    def test_add_sink(self, *, agent, agent2):
        agent.add_sink(agent2)
        assert agent2 in agent._sinks
        agent.add_sink(agent2)

    def test_channel_iterator(self, *, agent):
        agent.channel = Mock(name='channel', autospec=Channel)
        agent._channel_iterator = None
        it = agent.channel_iterator

        agent.channel.clone.assert_called_once_with(is_iterator=True)
        assert it is agent.channel.clone()
        agent.channel_iterator = [42]
        assert agent.channel_iterator == [42]

    def test_service(self, *, agent):
        with patch('faust.agents.agent.AgentService') as AgentService:
            service = agent._service
            AgentService.assert_called_once_with(
                agent,
                beacon=agent.app.beacon,
                loop=agent.app.loop,
            )
            assert service is AgentService()

    def test_label(self, *, agent):
        assert label(agent)
예제 #6
0
파일: test_agent.py 프로젝트: taybin/faust
class Test_Agent:
    @pytest.fixture
    def agent(self, *, app):
        @app.agent()
        async def myagent(stream):
            async for value in stream:
                yield value

        return myagent

    @pytest.fixture
    def isolated_agent(self, *, app):
        @app.agent(isolated_partitions=True)
        async def isoagent(stream):
            async for value in stream:
                yield value

        return isoagent

    @pytest.fixture
    def foo_topic(self, *, app):
        return app.topic("foo")

    @pytest.fixture
    def agent2(self, *, app, foo_topic):
        @app.agent(foo_topic)
        async def other_agent(stream):
            async for value in stream:
                value

        return other_agent

    def test_init_schema_and_channel(self, *, app):
        with pytest.raises(AssertionError):

            @app.agent(app.topic("foo"), schema=faust.Schema(key_type=bytes))
            async def foo():
                ...

    def test_init_key_type_and_channel(self, *, app):
        with pytest.raises(AssertionError):

            @app.agent(app.topic("foo"), key_type=bytes)
            async def foo():
                ...

    def test_init_value_type_and_channel(self, *, app):
        with pytest.raises(AssertionError):

            @app.agent(app.topic("foo"), value_type=bytes)
            async def foo():
                ...

    def test_isolated_partitions_cannot_have_concurrency(self, *, app):
        with pytest.raises(ImproperlyConfigured):

            @app.agent(isolated_partitions=True, concurrency=100)
            async def foo():
                ...

    def test_agent_call_reuse_stream(self, *, agent, app):
        stream = app.stream("foo")
        stream.concurrency_index = 1
        stream.active_partitions = {1, 2}
        actor = agent(stream=stream, index=1, active_partitions={1, 2})
        assert actor.stream is stream

    def test_cancel(self, *, agent):
        actor1 = Mock(name="actor1")
        actor2 = Mock(name="actor2")
        agent._actors = [actor1, actor2]
        agent.cancel()
        actor1.cancel.assert_called_once_with()
        actor2.cancel.assert_called_once_with()

    @pytest.mark.asyncio
    async def test_on_partitions_revoked(self, *, agent):
        revoked = {TP("foo", 0)}
        agent.on_shared_partitions_revoked = AsyncMock(name="ospr")
        await agent.on_partitions_revoked(revoked)
        agent.on_shared_partitions_revoked.assert_called_once_with(revoked)

    @pytest.mark.asyncio
    async def test_on_partitions_revoked__isolated(self, *, isolated_agent):
        revoked = {TP("foo", 0)}
        i = isolated_agent.on_isolated_partitions_revoked = AsyncMock(name="i")
        await isolated_agent.on_partitions_revoked(revoked)
        i.assert_called_once_with(revoked)

    @pytest.mark.asyncio
    async def test_on_partitions_assigned(self, *, agent):
        assigned = {TP("foo", 0)}
        agent.on_shared_partitions_assigned = AsyncMock(name="ospr")
        await agent.on_partitions_assigned(assigned)
        agent.on_shared_partitions_assigned.assert_called_once_with(assigned)

    @pytest.mark.asyncio
    async def test_on_partitions_assigned__isolated(self, *, isolated_agent):
        assigned = {TP("foo", 0)}
        i = isolated_agent.on_isolated_partitions_assigned = AsyncMock()
        await isolated_agent.on_partitions_assigned(assigned)
        i.assert_called_once_with(assigned)

    @pytest.mark.asyncio
    async def test_on_isolated_partitions_revoked(self, *, agent):
        tp = TP("foo", 0)
        aref = Mock(
            name="aref",
            autospec=Actor,
            on_isolated_partition_revoked=AsyncMock(),
        )
        agent._actor_by_partition = {tp: aref}

        await agent.on_isolated_partitions_revoked({tp})
        aref.on_isolated_partition_revoked.assert_called_once_with(tp)
        assert not agent._actor_by_partition
        await agent.on_isolated_partitions_revoked({tp})

    @pytest.mark.asyncio
    async def test_on_isolated_partitions_assigned(self, *, agent):
        agent._assign_isolated_partition = AsyncMock(name="aip")
        await agent.on_isolated_partitions_assigned({TP("foo", 0)})
        agent._assign_isolated_partition.assert_called_once_with(TP("foo", 0))

    @pytest.mark.asyncio
    async def test_assign_isolated_partition(self, *, agent):
        agent._on_first_isolated_partition_assigned = Mock(name="ofipa")
        agent._maybe_start_isolated = AsyncMock(name="maybe_start_isolated")
        agent._first_assignment_done = True

        tp = TP("foo", 606)
        await agent._assign_isolated_partition(tp)

        agent._on_first_isolated_partition_assigned.assert_not_called()
        agent._maybe_start_isolated.assert_called_once_with(tp)

        agent._first_assignment_done = False
        agent._actor_by_partition = set()
        await agent._assign_isolated_partition(tp)
        agent._on_first_isolated_partition_assigned.assert_called_once_with(tp)

    def test_on_first_isolated_partition_assigned(self, *, agent):
        aref = Mock(name="actor", autospec=Actor)
        agent._actors = [aref]
        agent._pending_active_partitions = set()
        tp = TP("foo", 303)
        agent._on_first_isolated_partition_assigned(tp)
        assert agent._actor_by_partition[tp] is aref
        assert agent._pending_active_partitions == {tp}
        agent._pending_active_partitions = None
        agent._on_first_isolated_partition_assigned(tp)

    @pytest.mark.asyncio
    async def test_maybe_start_isolated(self, *, isolated_agent):
        aref = Mock(
            name="actor",
            autospec=Actor,
            on_isolated_partition_assigned=AsyncMock(),
        )
        isolated_agent._start_isolated = AsyncMock(
            name="_start_isolated",
            return_value=aref,
        )
        tp = TP("foo", 303)
        await isolated_agent._maybe_start_isolated(tp)

        isolated_agent._start_isolated.assert_called_once_with(tp)
        assert isolated_agent._actor_by_partition[tp] is aref
        aref.on_isolated_partition_assigned.assert_called_once_with(tp)

    @pytest.mark.asyncio
    async def test_start_isolated(self, *, agent):
        agent._start_for_partitions = AsyncMock(
            name="agent._start_for_partitions", )
        ret = await agent._start_isolated(TP("foo", 0))
        agent._start_for_partitions.assert_called_once_with({TP("foo", 0)})
        assert ret is agent._start_for_partitions.coro()

    @pytest.mark.asyncio
    async def test_on_shared_partitions_revoked(self, *, agent):
        await agent.on_shared_partitions_revoked(set())

    @pytest.mark.asyncio
    async def test_on_shared_partitions_assigned(self, *, agent):
        await agent.on_shared_partitions_assigned(set())

    def test_info(self, *, agent):
        assert agent.info() == {
            "app": agent.app,
            "fun": agent.fun,
            "name": agent.name,
            "channel": agent.channel,
            "concurrency": agent.concurrency,
            "help": agent.help,
            "sink": agent._sinks,
            "on_error": agent._on_error,
            "supervisor_strategy": agent.supervisor_strategy,
            "isolated_partitions": agent.isolated_partitions,
        }

    def test_clone(self, *, agent):
        assert agent.clone(isolated_partitions=True).isolated_partitions

    def test_stream__active_partitions(self, *, agent):
        assert agent.stream(active_partitions={TP("foo", 0)})

    @pytest.mark.parametrize(
        "input,expected",
        [
            (ReqRepRequest("value", "reply_to", "correlation_id"), "value"),
            ("value", "value"),
        ],
    )
    def test_maybe_unwrap_reply_request(self, input, expected, *, agent):
        assert agent._maybe_unwrap_reply_request(input) == expected

    @pytest.mark.asyncio
    async def test_start_task(self, *, agent):
        agent._prepare_actor = AsyncMock(name="_prepare_actor")
        ret = await agent._start_task(index=0)
        agent._prepare_actor.assert_called_once_with(ANY, agent.beacon)
        assert ret is agent._prepare_actor.coro()

    @pytest.mark.asyncio
    async def test_prepare_actor__AsyncIterable(self, *, agent):
        aref = agent(index=0, active_partitions=None)
        with patch("asyncio.Task") as Task:
            agent._slurp = Mock(name="_slurp")
            agent._execute_actor = Mock(name="_execute_actor")
            beacon = Mock(name="beacon", autospec=Node)
            ret = await agent._prepare_actor(aref, beacon)
            agent._slurp.assert_called()
            coro = agent._slurp()
            agent._execute_actor.assert_called_once_with(coro, aref)
            Task.assert_called_once_with(agent._execute_actor(),
                                         loop=agent.loop)
            task = Task()
            assert task._beacon is beacon
            assert aref.actor_task is task
            assert aref in agent._actors
            assert ret is aref

    @pytest.mark.asyncio
    async def test_prepare_actor__Awaitable(self, *, agent2):
        aref = agent2(index=0, active_partitions=None)
        asyncio.ensure_future(aref.it).cancel()  # silence warning
        return
        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_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)
        asyncio.ensure_future(aref.it).cancel()  # silence warning
        agent2._sinks = [agent2]
        with pytest.raises(ImproperlyConfigured):
            await agent2._prepare_actor(
                aref,
                Mock(name="beacon", autospec=Node),
            )

    @pytest.mark.asyncio
    async def test_execute_actor(self, *, agent):
        coro = done_future()
        await agent._execute_actor(coro, Mock(name="aref", autospec=Actor))

    @pytest.mark.asyncio
    async def test_execute_actor__cancelled_stopped(self, *, agent):
        coro = FutureMock()
        coro.side_effect = asyncio.CancelledError()
        await agent.stop()
        with pytest.raises(asyncio.CancelledError):
            await agent._execute_actor(coro, Mock(name="aref", autospec=Actor))
        coro.assert_awaited()

    @pytest.mark.skip(reason="Fix is TBD")
    @pytest.mark.asyncio
    async def test_execute_actor__cancelled_running(self, *, agent):
        agent._on_error = AsyncMock(name="on_error")
        agent.log = Mock(name="log", autospec=CompositeLogger)
        aref = Mock(
            name="aref",
            autospec=Actor,
            crash=AsyncMock(),
        )
        agent.supervisor = Mock(name="supervisor")
        coro = FutureMock()
        exc = coro.side_effect = asyncio.CancelledError()
        await agent._execute_actor(coro, aref)
        coro.assert_awaited()

        aref.crash.assert_called_once_with(exc)
        agent.supervisor.wakeup.assert_called_once_with()
        agent._on_error.assert_not_called()

        agent._on_error = None
        await agent._execute_actor(coro, aref)

    @pytest.mark.asyncio
    async def test_execute_actor__raising(self, *, agent):
        agent._on_error = AsyncMock(name="on_error")
        agent.log = Mock(name="log", autospec=CompositeLogger)
        aref = Mock(
            name="aref",
            autospec=Actor,
            crash=AsyncMock(),
        )
        agent.supervisor = Mock(name="supervisor")
        coro = FutureMock()
        exc = coro.side_effect = KeyError("bar")
        await agent._execute_actor(coro, aref)
        coro.assert_awaited()

        aref.crash.assert_called_once_with(exc)
        agent.supervisor.wakeup.assert_called_once_with()
        agent._on_error.assert_called_once_with(agent, exc)

        agent._on_error = None
        await agent._execute_actor(coro, aref)

    @pytest.mark.asyncio
    async def test_slurp(self, *, agent, app):
        aref = agent(index=None, active_partitions=None)
        stream = aref.stream.get_active_stream()
        agent._delegate_to_sinks = AsyncMock(name="_delegate_to_sinks")
        agent._reply = AsyncMock(name="_reply")

        def on_delegate(value):
            raise StopAsyncIteration()

        word = Word("word")
        word_req = ReqRepRequest(word, "reply_to", "correlation_id")
        message1 = Mock(name="message1", autospec=Message)
        message2 = Mock(name="message2", autospec=Message)
        event1 = Event(app, None, word_req, {}, message1)
        event2 = Event(app, "key", "bar", {}, message2)
        values = [
            (event1, word),
            (event2, "bar"),
        ]

        class AIT:
            async def __aiter__(self):
                for event, value in values:
                    stream.current_event = event
                    yield value

        it = aiter(AIT())
        await agent._slurp(aref, it)

        agent._reply.assert_called_once_with(None, word, word_req.reply_to,
                                             word_req.correlation_id)
        agent._delegate_to_sinks.coro.assert_has_calls([
            call(word),
            call("bar"),
        ])

    @pytest.mark.asyncio
    async def test_slurp__headers(self, *, agent, app):
        agent.use_reply_headers = True
        aref = agent(index=None, active_partitions=None)
        stream = aref.stream.get_active_stream()
        agent._delegate_to_sinks = AsyncMock(name="_delegate_to_sinks")
        agent._reply = AsyncMock(name="_reply")

        def on_delegate(value):
            raise StopAsyncIteration()

        word = Word("word")
        message1 = Mock(name="message1", autospec=Message)
        headers1 = message1.headers = {
            "Faust-Ag-ReplyTo": "reply_to",
            "Faust-Ag-CorrelationId": "correlation_id",
        }
        message2 = Mock(name="message2", autospec=Message)
        headers2 = message2.headers = {}
        event1 = Event(app, None, word, headers1, message1)
        event2 = Event(app, "key", "bar", headers2, message2)
        values = [
            (event1, word, True),
            (event1, word, False),
            (event2, "bar", True),
        ]

        class AIT:
            async def __aiter__(self):
                for event, value, set_cur_event in values:
                    if set_cur_event:
                        stream.current_event = event
                    else:
                        stream.current_event = None
                    yield value

        it = aiter(AIT())
        await agent._slurp(aref, it)

        agent._reply.assert_called_once_with(None, word, "reply_to",
                                             "correlation_id")
        agent._delegate_to_sinks.coro.assert_has_calls([
            call(word),
            call("bar"),
        ])

    @pytest.mark.asyncio
    async def test_delegate_to_sinks(self, *, agent, agent2, foo_topic):
        agent2.send = AsyncMock(name="agent2.send")
        foo_topic.send = AsyncMock(name="foo_topic.send")
        sink_callback = Mock(name="sink_callback")
        sink_callback2_mock = Mock(name="sink_callback2_mock")

        async def sink_callback2(value):
            return sink_callback2_mock(value)

        agent._sinks = [
            agent2,
            foo_topic,
            sink_callback,
            sink_callback2,
        ]

        value = Mock(name="value")
        await agent._delegate_to_sinks(value)

        agent2.send.assert_called_once_with(value=value)
        foo_topic.send.assert_called_once_with(value=value)
        sink_callback.assert_called_once_with(value)
        sink_callback2_mock.assert_called_once_with(value)

    @pytest.mark.asyncio
    async def test_reply(self, *, agent):
        agent.app = Mock(
            name="app",
            autospec=App,
            send=AsyncMock(),
        )
        req = ReqRepRequest("value", "reply_to", "correlation_id")
        await agent._reply("key", "reply", req.reply_to, req.correlation_id)
        agent.app.send.assert_called_once_with(
            req.reply_to,
            key=None,
            value=ReqRepResponse(
                key="key",
                value="reply",
                correlation_id=req.correlation_id,
            ),
        )

    @pytest.mark.asyncio
    async def test_cast(self, *, agent):
        agent.send = AsyncMock(name="send")
        await agent.cast("value", key="key", partition=303)
        agent.send.assert_called_once_with(
            key="key",
            value="value",
            partition=303,
            headers=None,
            timestamp=None,
        )

    @pytest.mark.asyncio
    async def test_ask(self, *, agent):
        agent.app = Mock(
            name="app",
            autospec=App,
            maybe_start_client=AsyncMock(),
            _reply_consumer=Mock(
                autospec=ReplyConsumer,
                add=AsyncMock(),
            ),
        )
        pp = done_future()
        agent.ask_nowait = Mock(name="ask_nowait")
        agent.ask_nowait.return_value = done_future(pp)
        pp.correlation_id = "foo"

        await agent.ask(
            value="val",
            key="key",
            partition=303,
            correlation_id="correlation_id",
            headers={"k1": "v1"},
        )
        agent.ask_nowait.assert_called_once_with(
            "val",
            key="key",
            partition=303,
            reply_to=agent.app.conf.reply_to,
            correlation_id="correlation_id",
            force=True,
            timestamp=None,
            headers={"k1": "v1"},
        )
        agent.app._reply_consumer.add.assert_called_once_with(
            pp.correlation_id, pp)

    @pytest.mark.asyncio
    async def test_ask_nowait(self, *, agent):
        agent._create_req = Mock(name="_create_req")
        agent._create_req.return_value = ["V", None]
        agent.channel.send = AsyncMock(name="channel.send")
        res = await agent.ask_nowait(
            value="value",
            key="key",
            partition=303,
            timestamp=None,
            headers=None,
            reply_to="reply_to",
            correlation_id="correlation_id",
            force=True,
        )

        agent._create_req.assert_called_once_with("key", "value", "reply_to",
                                                  "correlation_id", None)
        agent.channel.send.assert_called_once_with(
            key="key",
            value=agent._create_req()[0],
            partition=303,
            timestamp=None,
            force=True,
            headers=None,
        )

        assert res.reply_to
        assert res.correlation_id

    @pytest.mark.asyncio
    async def test_ask_nowait__missing_reply_to(self, *, agent):
        with pytest.raises(TypeError):
            await agent.ask_nowait(
                value="value",
                key="key",
                partition=3034,
                reply_to=None,
            )

    def test_create_req(self, *, agent):
        agent.use_reply_headers = False
        agent._get_strtopic = Mock(name="_get_strtopic")
        with patch("faust.agents.agent.uuid4") as uuid4:
            uuid4.return_value = "vvv"
            reqrep = agent._create_req(key=b"key",
                                       value=b"value",
                                       reply_to="reply_to",
                                       headers={"k": "v"})[0]

            agent._get_strtopic.assert_called_once_with("reply_to")

            assert reqrep.value == b"value"
            assert reqrep.reply_to == agent._get_strtopic()
            assert reqrep.correlation_id == "vvv"

    def test_create_req__use_reply_headers(self, *, agent):
        agent.use_reply_headers = True
        agent._get_strtopic = Mock(name="_get_strtopic")
        with patch("faust.agents.agent.uuid4") as uuid4:
            uuid4.return_value = "vvv"
            value, h = agent._create_req(key=b"key",
                                         value=b"value",
                                         reply_to="reply_to",
                                         headers={"k": "v"})

            agent._get_strtopic.assert_called_once_with("reply_to")

            assert value == b"value"
            assert h["Faust-Ag-ReplyTo"] == agent._get_strtopic()
            assert h["Faust-Ag-CorrelationId"] == "vvv".encode()

    def test_create_req__model(self, *, agent):
        agent.use_reply_headers = False
        agent._get_strtopic = Mock(name="_get_strtopic")
        with patch("faust.agents.agent.uuid4") as uuid4:
            uuid4.return_value = "vvv"
            value = Word("foo")
            reqrep = agent._create_req(key=b"key",
                                       value=value,
                                       reply_to="reply_to",
                                       headers={"h1": "h2"})[0]
            assert isinstance(reqrep, ReqRepRequest)

            agent._get_strtopic.assert_called_once_with("reply_to")
            assert isinstance(reqrep, ModelReqRepRequest)

            assert reqrep.value is value
            assert reqrep.reply_to == agent._get_strtopic()
            assert reqrep.correlation_id == "vvv"

    def test_create_req__requires_reply_to(self, *, agent):
        with pytest.raises(TypeError):
            agent._create_req(
                key=b"key",
                value=b"value",
                reply_to=None,
            )

    @pytest.mark.parametrize(
        "value,expected_class",
        [
            (b"value", ReqRepResponse),
            (Word("foo"), ModelReqRepResponse),
        ],
    )
    def test_response_class(self, value, expected_class, *, agent):
        assert agent._response_class(value) is expected_class

    @pytest.mark.asyncio
    async def test_send(self, *, agent):
        agent.channel = Mock(
            name="channel",
            autospec=Channel,
            send=AsyncMock(),
        )
        agent._create_req = Mock(name="_create_req")
        agent._create_req.return_value = ("V", {"k2": "v2"})
        callback = Mock(name="callback")

        ret = await agent.send(
            key=b"key",
            value=b"value",
            partition=303,
            timestamp=None,
            headers={"k": "v"},
            key_serializer="raw",
            value_serializer="raw",
            callback=callback,
            reply_to="reply_to",
            correlation_id="correlation_id",
            force=True,
        )

        agent._create_req.assert_called_once_with(
            b"key",
            b"value",
            "reply_to",
            "correlation_id",
            {"k": "v"},
        )

        agent.channel.send.assert_called_once_with(
            key=b"key",
            value=agent._create_req()[0],
            partition=303,
            timestamp=None,
            headers={"k2": "v2"},
            key_serializer="raw",
            value_serializer="raw",
            force=True,
        )

        assert ret is agent.channel.send.coro()

    @pytest.mark.asyncio
    async def test_send__without_reply_to(self, *, agent):
        agent.channel = Mock(
            name="channel",
            autospec=Channel,
            send=AsyncMock(),
        )
        agent._create_req = Mock(name="_create_req")
        callback = Mock(name="callback")

        ret = await agent.send(
            key=b"key",
            value=b"value",
            partition=303,
            timestamp=None,
            headers={"k": "v"},
            key_serializer="raw",
            value_serializer="raw",
            callback=callback,
            reply_to=None,
            correlation_id="correlation_id",
            force=True,
        )

        agent._create_req.assert_not_called()

        agent.channel.send.assert_called_once_with(
            key=b"key",
            value=b"value",
            partition=303,
            timestamp=None,
            headers={"k": "v"},
            key_serializer="raw",
            value_serializer="raw",
            force=True,
        )

        assert ret is agent.channel.send.coro()

    def test_get_strtopic__agent(self, *, agent, agent2):
        assert agent._get_strtopic(agent2) == agent2.channel.get_topic_name()

    def test_get_strtopic__topic(self, *, agent, foo_topic):
        assert agent._get_strtopic(foo_topic) == foo_topic.get_topic_name()

    def test_get_strtopic__str(self, *, agent):
        assert agent._get_strtopic("bar") == "bar"

    def test_get_strtopic__channel_raises(self, *, agent, app):
        with pytest.raises(ValueError):
            agent._get_strtopic(app.channel())

    def test_get_topic_names(self, *, agent, app):
        agent.channel = app.topic("foo")
        assert agent.get_topic_names() == ("foo", )

    def test_get_topic_names__channel(self, *, agent, app):
        agent.channel = app.channel()
        assert agent.get_topic_names() == []

    def test_repr(self, *, agent):
        assert repr(agent)

    def test_channel(self, *, agent):
        agent._prepare_channel = Mock(name="_prepare_channel")
        agent._channel = None
        channel = agent.channel
        agent._prepare_channel.assert_called_once_with(
            agent._channel_arg,
            schema=agent._schema,
            key_type=agent._key_type,
            value_type=agent._value_type,
            **agent._channel_kwargs,
        )
        assert channel is agent._prepare_channel.return_value
        assert agent._channel is channel

    def test_prepare_channel__not_channel(self, *, agent):
        with pytest.raises(TypeError):
            agent._prepare_channel(object())

    def test_add_sink(self, *, agent, agent2):
        agent.add_sink(agent2)
        assert agent2 in agent._sinks
        agent.add_sink(agent2)

    def test_channel_iterator(self, *, agent):
        agent.channel = Mock(name="channel", autospec=Channel)
        agent._channel_iterator = None
        it = agent.channel_iterator

        agent.channel.clone.assert_called_once_with(is_iterator=False)
        assert it is agent.channel.clone()
        agent.channel_iterator = [42]
        assert agent.channel_iterator == [42]

    def test_label(self, *, agent):
        assert label(agent)

    async def test_context_calls_sink(self, *, agent):
        class SinkCalledException(Exception):
            pass

        def dummy_sink(_):
            raise SinkCalledException()

        agent.add_sink(dummy_sink)
        async with agent.test_context() as agent_mock:
            with pytest.raises(SinkCalledException):
                await agent_mock.put("hello")