51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from cmvr_edge_ai.transports.grpc import (
|
||
|
|
BidiRequestStream,
|
||
|
|
BidiRequestStreamClosed,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_close_wakes_blocked_sender_without_accepting_its_request() -> None:
|
||
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
||
|
|
stream = BidiRequestStream[str](capacity=1)
|
||
|
|
await stream.send("accepted")
|
||
|
|
blocked_sender = asyncio.create_task(stream.send("blocked"))
|
||
|
|
await asyncio.sleep(0)
|
||
|
|
assert blocked_sender.done() is False
|
||
|
|
|
||
|
|
await stream.close()
|
||
|
|
with pytest.raises(BidiRequestStreamClosed, match="closed"):
|
||
|
|
await blocked_sender
|
||
|
|
drained = [request async for request in stream]
|
||
|
|
return drained, stream.stats()
|
||
|
|
|
||
|
|
drained, stats = asyncio.run(exercise())
|
||
|
|
|
||
|
|
assert drained == ["accepted"]
|
||
|
|
assert stats.enqueued == 1
|
||
|
|
assert stats.blocked_puts == 1
|
||
|
|
assert stats.closed is True
|
||
|
|
|
||
|
|
|
||
|
|
def test_bidirectional_request_stream_allows_exactly_one_consumer() -> None:
|
||
|
|
async def exercise() -> None:
|
||
|
|
stream = BidiRequestStream[str]()
|
||
|
|
await stream.send("first")
|
||
|
|
first_consumer = stream.__aiter__()
|
||
|
|
assert await anext(first_consumer) == "first"
|
||
|
|
|
||
|
|
second_consumer = stream.__aiter__()
|
||
|
|
with pytest.raises(RuntimeError, match="exactly one consumer"):
|
||
|
|
await anext(second_consumer)
|
||
|
|
|
||
|
|
await stream.close()
|
||
|
|
with pytest.raises(StopAsyncIteration):
|
||
|
|
await anext(first_consumer)
|
||
|
|
|
||
|
|
asyncio.run(exercise())
|