async def test_message_handler(self):
side_effect = None
async def raw_func(data):
nonlocal side_effect
assert isinstance(data, dict)
side_effect = True
handler = self.factory(raw_func)
await handler(record_factory(), None)
assert side_effect is True
async def test_message_handler_map_types(self):
class Foo(pydantic.BaseModel):
foo: str
async def handle_func(ob: Foo, schema, record, app, subscriber, span: opentracing.Span):
assert ob.foo == "bar"
assert schema == "Foo:1"
assert record is not None
assert app is not None
assert subscriber is not None
assert span is not None
mh = self.factory(handle_func)
await mh.handle(record_factory(), MagicMock())
def test_publish_callback(self, metrics):
fut = Future()
fut.set_result(record_factory())
published_callback("topic", time.time() - 1, fut)
metrics["PUBLISHED_MESSAGES"].labels.assert_called_with(
stream_id="topic", partition=0, error="none"
)
metrics["PUBLISHED_MESSAGES"].labels().inc()
metrics["PRODUCER_TOPIC_OFFSET"].labels.assert_called_with(stream_id="topic", partition=0)
metrics["PRODUCER_TOPIC_OFFSET"].labels().set.assert_called_with(0)
metrics["PUBLISHED_MESSAGES_TIME"].labels.assert_called_with(stream_id="topic")
assert metrics["PUBLISHED_MESSAGES_TIME"].labels().observe.mock_calls[0].args[
0
] == pytest.approx(1, 0.1)
async def test_run_exits_when_fut_closed_fut(self, subscription):
sub = subscription
consumer = AsyncMock()
consumer.getmany.return_value = {
"": [record_factory() for _ in range(10)]
}
sub._consumer = consumer
sub._running = True
async def _handle_message(record):
await asyncio.sleep(0.03)
with patch.object(sub, "_handler", _handle_message):
task = asyncio.create_task(sub._consume())
await asyncio.sleep(0.01)
stop_task = asyncio.create_task(sub.stop())
await asyncio.sleep(0.01)
sub._close.set_result(None)
await asyncio.wait([stop_task, task])
async def test_run_exits_when_fut_closed_fut(self):
sub = SubscriptionConsumer(
Application(),
Subscription("foo", lambda record: 1, "group"),
)
consumer = AsyncMock()
consumer.getmany.return_value = {"": [record_factory() for _ in range(10)]}
sub._consumer = consumer
sub._running = True
async def _handle_message(record):
await asyncio.sleep(0.03)
with patch.object(sub, "_handle_message", _handle_message):
task = asyncio.create_task(sub._run())
await asyncio.sleep(0.01)
stop_task = asyncio.create_task(sub.stop())
await asyncio.sleep(0.01)
sub._close_fut.set_result(None)
await asyncio.wait([stop_task, task])
async def test_message_handler(self):
async def raw_func(data):
assert isinstance(data, dict)
mh = self.factory(raw_func)
await mh.handle(record_factory(), None)
async def _publish():
return record_factory()
