def test_add_remove():
batcher = FlowControlBatcher()
batcher.add(FlowControlRequest(allowed_bytes=10, allowed_messages=3))
restart_1 = batcher.request_for_restart()
assert restart_1.allowed_bytes == 10
assert restart_1.allowed_messages == 3
batcher.on_messages(
[SequencedMessage(size_bytes=2),
SequencedMessage(size_bytes=3)])
restart_2 = batcher.request_for_restart()
assert restart_2.allowed_bytes == 5
assert restart_2.allowed_messages == 1
def test_restart_clears_send():
batcher = FlowControlBatcher()
batcher.add(FlowControlRequest(allowed_bytes=10, allowed_messages=3))
assert batcher.should_expedite()
to_send = batcher.release_pending_request()
assert to_send.allowed_bytes == 10
assert to_send.allowed_messages == 3
restart_1 = batcher.request_for_restart()
assert restart_1.allowed_bytes == 10
assert restart_1.allowed_messages == 3
assert not batcher.should_expedite()
assert batcher.release_pending_request() is None
def __init__(
self,
initial: InitialSubscribeRequest,
token_flush_seconds: float,
factory: ConnectionFactory[SubscribeRequest, SubscribeResponse],
):
self._initial = initial
self._token_flush_seconds = token_flush_seconds
self._connection = RetryingConnection(factory, self)
self._outstanding_flow_control = FlowControlBatcher()
self._reinitializing = False
self._last_received_offset = None
self._message_queue = asyncio.Queue()
self._receiver = None
self._flusher = None