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)
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)
async def test_send_and_wait(self, producer):
producer.send = AsyncMock(return_value=done_future(done_future()))
await producer.send_and_wait(
'topic', 'k', 'v', 3, 100, [('a', 'b')],
transactional_id='tid')
producer.send.assert_called_once_with(
'topic',
key='k',
value='v',
partition=3,
timestamp=100,
headers=[('a', 'b')],
transactional_id='tid',
)
async def test_already_done(self):
async def wrapped():
pass
x = StampedeWrapper(wrapped)
x.fut = done_future('foo')
assert await x() == 'foo'
async def send(self, topic: str, key: Optional[bytes],
value: Optional[bytes],
partition: Optional[int],
timestamp: Optional[float],
headers: Optional[HeadersArg],
*,
transactional_id: str = None) -> Awaitable[RecordMetadata]:
res = await self.send_and_wait(
topic, key, value, partition, timestamp, headers)
return cast(Awaitable[RecordMetadata], done_future(res))
async def test_throw__notify_pending_waiters(self):
flow_control = FlowControlEvent(initially_suspended=False)
queue = ThrowableQueue(flow_control=flow_control, maxsize=1)
raised = 0
async def waiter():
try:
await queue.get()
except KeyError:
nonlocal raised
raised += 1
queue._getters.append(done_future())
fut = asyncio.ensure_future(waiter())
await asyncio.sleep(0.01)
await queue.throw(KeyError())
await asyncio.gather(fut)
assert raised == 1
async def send(*args, **kwargs):
on_send(*args, **kwargs)
return done_future()
async def _call(*args, **kwargs):
consumer._method_queue._call(*args, **kwargs)
return done_future()
async def test_execute_task(self, *, agent):
coro = done_future()
await agent._execute_task(coro, Mock(name='aref', autospec=Actor))
async def test_wait_for_subscription(self, *, con):
con._subscription_done = None
await con.wait_for_subscriptions()
con._subscription_done = done_future()
await con.wait_for_subscriptions()
async def send(self, topic: str, key: Optional[bytes],
value: Optional[bytes],
partition: Optional[int]) -> Awaitable[RecordMetadata]:
res = await self.send_and_wait(topic, key, value, partition)
return cast(Awaitable[RecordMetadata], done_future(res))
async def test_maybe_wait_for_subscriptions(self, *, con):
con._subscription_done = done_future()
await con.maybe_wait_for_subscriptions()
async def test_wait_for_subscriptions__notset(self, *, con):
with patch('asyncio.Future') as Future:
Future.return_value = done_future()
await con.wait_for_subscriptions()
async def on_call(*args, **kwargs):
return done_future('value')
async def test_execute_actor(self, *, agent):
coro = done_future()
await agent._execute_actor(coro, Mock(name="aref", autospec=Actor))
