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'),
])
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"),
])
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,
),
)
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,
),
)
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)
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")
