async def test_basic_commit_after_timeout(
committer: Committer,
default_connection,
initial_request,
asyncio_sleep,
sleep_queues,
):
sleep_called = sleep_queues[FLUSH_SECONDS].called
sleep_results = sleep_queues[FLUSH_SECONDS].results
cursor1 = Cursor(offset=321)
cursor2 = Cursor(offset=1)
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue
)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue
)
read_result_queue.put_nowait(StreamingCommitCursorResponse(initial={}))
write_result_queue.put_nowait(None)
async with committer:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Commit cursors
commit_fut1 = asyncio.ensure_future(committer.commit(cursor1))
commit_fut2 = asyncio.ensure_future(committer.commit(cursor2))
empty_fut = asyncio.ensure_future(committer.wait_until_empty())
assert not commit_fut1.done()
assert not commit_fut2.done()
assert not empty_fut.done()
# Wait for writes to be waiting
await sleep_called.get()
asyncio_sleep.assert_called_with(FLUSH_SECONDS)
# Handle the connection write
await sleep_results.put(None)
await write_called_queue.get()
await write_result_queue.put(None)
# Called with second cursor
default_connection.write.assert_has_calls(
[call(initial_request), call(as_request(cursor2))]
)
assert not commit_fut1.done()
assert not commit_fut2.done()
assert not empty_fut.done()
# Send the connection response with 1 ack since only one request was sent.
await read_result_queue.put(as_response(count=1))
await commit_fut1
await commit_fut2
await empty_fut
async def test_out_of_order_receipt_failure(
subscriber: Subscriber,
default_connection,
initial_request,
asyncio_sleep,
sleep_queues,
):
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
flow = FlowControlRequest(allowed_messages=100, allowed_bytes=100)
message_1 = SequencedMessage(cursor=Cursor(offset=3), size_bytes=5)
message_2 = SequencedMessage(cursor=Cursor(offset=5), size_bytes=10)
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue
)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue
)
read_result_queue.put_nowait(SubscribeResponse(initial={}))
write_result_queue.put_nowait(None)
async with subscriber:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Send tokens.
flow_fut = asyncio.ensure_future(subscriber.allow_flow(flow))
assert not flow_fut.done()
# Handle the inline write since initial tokens are 100% of outstanding.
await write_called_queue.get()
await write_result_queue.put(None)
await flow_fut
default_connection.write.assert_has_calls(
[call(initial_request), call(as_request(flow))]
)
read_fut = asyncio.ensure_future(subscriber.read())
# Send out of order messages to the subscriber.
await read_result_queue.put(as_response([message_2, message_1]))
# Wait for the next read call
await read_called_queue.get()
try:
await read_fut
assert False
except GoogleAPICallError as e:
assert e.grpc_status_code == StatusCode.FAILED_PRECONDITION
pass
async def test_ack(subscriber: AsyncSingleSubscriber, underlying, transformer,
ack_set_tracker):
ack_called_queue = asyncio.Queue()
ack_result_queue = asyncio.Queue()
ack_set_tracker.ack.side_effect = make_queue_waiter(
ack_called_queue, ack_result_queue)
async with subscriber:
message_1 = SequencedMessage(cursor=Cursor(offset=1), size_bytes=5)
message_2 = SequencedMessage(cursor=Cursor(offset=2), size_bytes=10)
underlying.read.return_value = message_1
read_1: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1)])
assert read_1.message_id == "1"
underlying.read.return_value = message_2
read_2: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1), call(2)])
assert read_2.message_id == "2"
read_2.ack()
await ack_called_queue.get()
await ack_result_queue.put(None)
ack_set_tracker.ack.assert_has_calls([call(2)])
read_1.ack()
await ack_called_queue.get()
await ack_result_queue.put(None)
ack_set_tracker.ack.assert_has_calls([call(2), call(1)])
async def test_ack_failure(
subscriber: SinglePartitionSingleSubscriber,
underlying,
transformer,
ack_set_tracker,
):
ack_called_queue = asyncio.Queue()
ack_result_queue = asyncio.Queue()
ack_set_tracker.ack.side_effect = make_queue_waiter(
ack_called_queue, ack_result_queue)
async with subscriber:
message = SequencedMessage(cursor=Cursor(offset=1), size_bytes=5)
underlying.read.return_value = message
read: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1)])
read.ack()
await ack_called_queue.get()
ack_set_tracker.ack.assert_has_calls([call(1)])
await ack_result_queue.put(FailedPrecondition("Bad ack"))
async def sleep_forever():
await asyncio.sleep(float("inf"))
underlying.read.side_effect = sleep_forever
with pytest.raises(FailedPrecondition):
await subscriber.read()
async def test_nack_calls_ack(
subscriber: SinglePartitionSingleSubscriber,
underlying,
transformer,
ack_set_tracker,
nack_handler,
):
ack_called_queue = asyncio.Queue()
ack_result_queue = asyncio.Queue()
ack_set_tracker.ack.side_effect = make_queue_waiter(
ack_called_queue, ack_result_queue)
async with subscriber:
message = SequencedMessage(cursor=Cursor(offset=1), size_bytes=5)
underlying.read.return_value = message
read: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1)])
def on_nack(nacked: PubsubMessage, ack: Callable[[], None]):
assert nacked.message_id == "1"
ack()
nack_handler.on_nack.side_effect = on_nack
read.nack()
await ack_called_queue.get()
await ack_result_queue.put(None)
ack_set_tracker.ack.assert_has_calls([call(1)])
async def test_basic_publish_after_timeout(
publisher: Publisher,
default_connection,
initial_request,
asyncio_sleep,
sleep_queues,
):
sleep_called = sleep_queues[FLUSH_SECONDS].called
sleep_results = sleep_queues[FLUSH_SECONDS].results
message1 = PubSubMessage(data=b"abc")
message2 = PubSubMessage(data=b"def")
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue)
read_result_queue.put_nowait(PublishResponse(initial_response={}))
async with publisher:
# Set up connection
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Write messages
publish_fut1 = asyncio.ensure_future(publisher.publish(message1))
publish_fut2 = asyncio.ensure_future(publisher.publish(message2))
assert not publish_fut1.done()
assert not publish_fut2.done()
# Wait for writes to be waiting
await sleep_called.get()
asyncio_sleep.assert_called_with(FLUSH_SECONDS)
# Handle the connection write
write_future = asyncio.Future()
async def write(val: PublishRequest):
write_future.set_result(None)
default_connection.write.side_effect = write
await sleep_results.put(None)
await write_future
default_connection.write.assert_has_calls([
call(initial_request),
call(as_publish_request([message1, message2]))
])
assert not publish_fut1.done()
assert not publish_fut2.done()
# Send the connection response
await read_result_queue.put(as_publish_response(100))
cursor1 = (await publish_fut1).cursor
cursor2 = (await publish_fut2).cursor
assert cursor1.offset == 100
assert cursor2.offset == 101
async def test_basic_assign(assigner: Assigner, default_connection,
initial_request):
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue)
write_result_queue.put_nowait(None)
async with assigner:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Wait for the first assignment
assign_fut1 = asyncio.ensure_future(assigner.get_assignment())
assert not assign_fut1.done()
partitions = {Partition(2), Partition(7)}
# Send the first assignment.
await read_result_queue.put(as_response(partitions=partitions))
assert (await assign_fut1) == partitions
# Get the next assignment: should send an ack on the stream
assign_fut2 = asyncio.ensure_future(assigner.get_assignment())
await write_called_queue.get()
await write_result_queue.put(None)
default_connection.write.assert_has_calls(
[call(initial_request), call(ack_request())])
partitions = {Partition(5)}
# Send the second assignment.
await read_called_queue.get()
await read_result_queue.put(as_response(partitions=partitions))
assert (await assign_fut2) == partitions
async def test_handle_reset(
subscriber: SinglePartitionSingleSubscriber,
underlying,
transformer,
ack_set_tracker,
):
ack_called_queue = asyncio.Queue()
ack_result_queue = asyncio.Queue()
ack_set_tracker.ack.side_effect = make_queue_waiter(
ack_called_queue, ack_result_queue
)
async with subscriber:
message_1 = SequencedMessage(cursor=Cursor(offset=1), size_bytes=5)
underlying.read.return_value = message_1
read_1: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1)])
assert read_1.message_id == "1"
assert read_1.ack_id == ack_id(0, 1)
await subscriber.handle_reset()
ack_set_tracker.clear_and_commit.assert_called_once()
# Message ACKed after reset. Its flow control tokens are refilled
# but offset not committed (verified below after message 2).
read_1.ack()
message_2 = SequencedMessage(cursor=Cursor(offset=2), size_bytes=10)
underlying.read.return_value = message_2
read_2: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1), call(2)])
assert read_2.message_id == "2"
assert read_2.ack_id == ack_id(1, 2)
read_2.ack()
await ack_called_queue.get()
await ack_result_queue.put(None)
underlying.allow_flow.assert_has_calls(
[
call(FlowControlRequest(allowed_messages=1000, allowed_bytes=1000,)),
call(FlowControlRequest(allowed_messages=1, allowed_bytes=5,)),
call(FlowControlRequest(allowed_messages=1, allowed_bytes=10,)),
]
)
ack_set_tracker.ack.assert_has_calls([call(2)])
async def test_publishes_retried_on_restart(
committer: Committer,
default_connection,
initial_request,
asyncio_sleep,
sleep_queues,
):
sleep_called = sleep_queues[FLUSH_SECONDS].called
sleep_results = sleep_queues[FLUSH_SECONDS].results
cursor1 = Cursor(offset=321)
cursor2 = Cursor(offset=1)
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue)
read_result_queue.put_nowait(StreamingCommitCursorResponse(initial={}))
write_result_queue.put_nowait(None)
async with committer:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Write message 1
commit_fut1 = asyncio.ensure_future(committer.commit(cursor1))
assert not commit_fut1.done()
# Wait for writes to be waiting
await sleep_called.get()
asyncio_sleep.assert_called_with(FLUSH_SECONDS)
# Handle the connection write
await sleep_results.put(None)
await write_called_queue.get()
await write_result_queue.put(None)
default_connection.write.assert_has_calls(
[call(initial_request),
call(as_request(cursor1))])
assert not commit_fut1.done()
# Wait for writes to be waiting
await sleep_called.get()
asyncio_sleep.assert_has_calls(
[call(FLUSH_SECONDS), call(FLUSH_SECONDS)])
# Write message 2
commit_fut2 = asyncio.ensure_future(committer.commit(cursor2))
assert not commit_fut2.done()
# Handle the connection write
await sleep_results.put(None)
await write_called_queue.get()
await write_result_queue.put(None)
default_connection.write.assert_has_calls([
call(initial_request),
call(as_request(cursor1)),
call(as_request(cursor2)),
])
assert not commit_fut1.done()
assert not commit_fut2.done()
# Fail the connection with a retryable error
await read_called_queue.get()
await read_result_queue.put(InternalServerError("retryable"))
await sleep_queues[_MIN_BACKOFF_SECS].called.get()
await sleep_queues[_MIN_BACKOFF_SECS].results.put(None)
# Reinitialization
await write_called_queue.get()
await write_result_queue.put(None)
await read_called_queue.get()
await read_result_queue.put(StreamingCommitCursorResponse(initial={}))
# Re-sending messages on the new stream
await write_called_queue.get()
await write_result_queue.put(None)
asyncio_sleep.assert_has_calls([
call(FLUSH_SECONDS),
call(FLUSH_SECONDS),
call(FLUSH_SECONDS),
call(_MIN_BACKOFF_SECS),
])
default_connection.write.assert_has_calls([
# Aggregates response calls on second pass
call(initial_request),
call(as_request(cursor2)),
])
# Sending the response for the one commit finishes both
await read_called_queue.get()
await read_result_queue.put(as_response(count=1))
await commit_fut1
await commit_fut2
async def test_commits_multi_cycle(
committer: Committer,
default_connection,
initial_request,
asyncio_sleep,
sleep_queues,
):
sleep_called = sleep_queues[FLUSH_SECONDS].called
sleep_results = sleep_queues[FLUSH_SECONDS].results
cursor1 = Cursor(offset=321)
cursor2 = Cursor(offset=1)
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue)
read_result_queue.put_nowait(StreamingCommitCursorResponse(initial={}))
write_result_queue.put_nowait(None)
async with committer:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Write message 1
commit_fut1 = asyncio.ensure_future(committer.commit(cursor1))
assert not commit_fut1.done()
# Wait for writes to be waiting
await sleep_called.get()
asyncio_sleep.assert_called_with(FLUSH_SECONDS)
# Handle the connection write
await sleep_results.put(None)
await write_called_queue.get()
await write_result_queue.put(None)
default_connection.write.assert_has_calls(
[call(initial_request),
call(as_request(cursor1))])
assert not commit_fut1.done()
# Wait for writes to be waiting
await sleep_called.get()
asyncio_sleep.assert_has_calls(
[call(FLUSH_SECONDS), call(FLUSH_SECONDS)])
# Write message 2
commit_fut2 = asyncio.ensure_future(committer.commit(cursor2))
assert not commit_fut2.done()
# Handle the connection write
await sleep_results.put(None)
await write_called_queue.get()
await write_result_queue.put(None)
default_connection.write.assert_has_calls([
call(initial_request),
call(as_request(cursor1)),
call(as_request(cursor2)),
])
assert not commit_fut1.done()
assert not commit_fut2.done()
# Send the connection responses
await read_result_queue.put(as_response(count=2))
await commit_fut1
await commit_fut2
async def test_basic_flow_control_after_timeout(
subscriber: Subscriber,
default_connection,
initial_request,
asyncio_sleep,
sleep_queues,
):
sleep_called = sleep_queues[FLUSH_SECONDS].called
sleep_results = sleep_queues[FLUSH_SECONDS].results
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
flow_1 = FlowControlRequest(allowed_messages=100, allowed_bytes=100)
flow_2 = FlowControlRequest(allowed_messages=5, allowed_bytes=10)
flow_3 = FlowControlRequest(allowed_messages=10, allowed_bytes=5)
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue
)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue
)
read_result_queue.put_nowait(SubscribeResponse(initial={}))
write_result_queue.put_nowait(None)
async with subscriber:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Send tokens.
flow_fut1 = asyncio.ensure_future(subscriber.allow_flow(flow_1))
assert not flow_fut1.done()
# Handle the inline write since initial tokens are 100% of outstanding.
await write_called_queue.get()
await write_result_queue.put(None)
await flow_fut1
default_connection.write.assert_has_calls(
[call(initial_request), call(as_request(flow_1))]
)
# Should complete without writing to the connection
await subscriber.allow_flow(flow_2)
await subscriber.allow_flow(flow_3)
# Wait for writes to be waiting
await sleep_called.get()
asyncio_sleep.assert_called_with(FLUSH_SECONDS)
# Handle the connection write
await sleep_results.put(None)
await write_called_queue.get()
await write_result_queue.put(None)
# Called with aggregate
default_connection.write.assert_has_calls(
[
call(initial_request),
call(as_request(flow_1)),
call(
as_request(
FlowControlRequest(allowed_messages=15, allowed_bytes=15)
)
),
]
)
async def test_publishes_retried_on_restart(
publisher: Publisher,
default_connection,
initial_request,
asyncio_sleep,
sleep_queues,
):
sleep_called = sleep_queues[FLUSH_SECONDS].called
sleep_results = sleep_queues[FLUSH_SECONDS].results
message1 = PubSubMessage(data=b"abc")
message2 = PubSubMessage(data=b"def")
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue)
write_result_queue.put_nowait(None)
read_result_queue.put_nowait(PublishResponse(initial_response={}))
async with publisher:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Write message 1
publish_fut1 = asyncio.ensure_future(publisher.publish(message1))
assert not publish_fut1.done()
# Wait for writes to be waiting
await sleep_called.get()
asyncio_sleep.assert_called_with(FLUSH_SECONDS)
# Handle the connection write
await sleep_results.put(None)
await write_called_queue.get()
await write_result_queue.put(None)
default_connection.write.assert_has_calls(
[call(initial_request),
call(as_publish_request([message1]))])
assert not publish_fut1.done()
# Wait for writes to be waiting
await sleep_called.get()
asyncio_sleep.assert_has_calls(
[call(FLUSH_SECONDS), call(FLUSH_SECONDS)])
# Write message 2
publish_fut2 = asyncio.ensure_future(publisher.publish(message2))
assert not publish_fut2.done()
# Handle the connection write
await sleep_results.put(None)
await write_called_queue.get()
await write_result_queue.put(None)
default_connection.write.assert_has_calls([
call(initial_request),
call(as_publish_request([message1])),
call(as_publish_request([message2])),
])
assert not publish_fut1.done()
assert not publish_fut2.done()
# Fail the connection with a retryable error
await read_called_queue.get()
await read_result_queue.put(InternalServerError("retryable"))
await sleep_queues[_MIN_BACKOFF_SECS].called.get()
await sleep_queues[_MIN_BACKOFF_SECS].results.put(None)
# Reinitialization
await write_called_queue.get()
write_result_queue.put_nowait(None)
await read_called_queue.get()
read_result_queue.put_nowait(PublishResponse(initial_response={}))
# Re-sending messages on the new stream
await write_called_queue.get()
await write_result_queue.put(None)
await write_called_queue.get()
await write_result_queue.put(None)
asyncio_sleep.assert_has_calls([
call(FLUSH_SECONDS),
call(FLUSH_SECONDS),
call(FLUSH_SECONDS),
call(_MIN_BACKOFF_SECS),
])
default_connection.write.assert_has_calls([
call(initial_request),
call(as_publish_request([message1])),
call(as_publish_request([message2])),
call(initial_request),
call(as_publish_request([message1])),
call(as_publish_request([message2])),
])
async def test_wait_until_empty_completes_on_failure(
committer: Committer,
default_connection,
initial_request,
asyncio_sleep,
sleep_queues,
):
sleep_called = sleep_queues[FLUSH_SECONDS].called
sleep_results = sleep_queues[FLUSH_SECONDS].results
cursor1 = Cursor(offset=1)
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue
)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue
)
read_result_queue.put_nowait(StreamingCommitCursorResponse(initial={}))
write_result_queue.put_nowait(None)
async with committer:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# New committer is empty.
await committer.wait_until_empty()
# Write message 1
commit_fut1 = asyncio.ensure_future(committer.commit(cursor1))
empty_fut = asyncio.ensure_future(committer.wait_until_empty())
assert not commit_fut1.done()
assert not empty_fut.done()
# Wait for writes to be waiting
await sleep_called.get()
asyncio_sleep.assert_called_with(FLUSH_SECONDS)
# Handle the connection write
await sleep_results.put(None)
await write_called_queue.get()
await write_result_queue.put(None)
default_connection.write.assert_has_calls(
[call(initial_request), call(as_request(cursor1))]
)
assert not commit_fut1.done()
assert not empty_fut.done()
# Wait for writes to be waiting
await sleep_called.get()
asyncio_sleep.assert_has_calls([call(FLUSH_SECONDS), call(FLUSH_SECONDS)])
# Fail the connection with a permanent error
await read_called_queue.get()
await read_result_queue.put(InvalidArgument("permanent"))
with pytest.raises(InvalidArgument):
await empty_fut
async def test_flow_resent_on_restart(
subscriber: Subscriber,
default_connection,
initial_request,
asyncio_sleep,
sleep_queues,
):
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
flow_1 = FlowControlRequest(allowed_messages=100, allowed_bytes=100)
flow_2 = FlowControlRequest(allowed_messages=5, allowed_bytes=10)
flow_3 = FlowControlRequest(allowed_messages=10, allowed_bytes=5)
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue
)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue
)
read_result_queue.put_nowait(SubscribeResponse(initial={}))
write_result_queue.put_nowait(None)
async with subscriber:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Send tokens.
flow_fut1 = asyncio.ensure_future(subscriber.allow_flow(flow_1))
assert not flow_fut1.done()
# Handle the inline write since initial tokens are 100% of outstanding.
await write_called_queue.get()
await write_result_queue.put(None)
await flow_fut1
default_connection.write.assert_has_calls(
[call(initial_request), call(as_request(flow_1))]
)
# Should complete without writing to the connection
await subscriber.allow_flow(flow_2)
await subscriber.allow_flow(flow_3)
# Fail the connection with a retryable error
await read_called_queue.get()
await read_result_queue.put(InternalServerError("retryable"))
await sleep_queues[_MIN_BACKOFF_SECS].called.get()
await sleep_queues[_MIN_BACKOFF_SECS].results.put(None)
# Reinitialization
await write_called_queue.get()
await write_result_queue.put(None)
await read_called_queue.get()
await read_result_queue.put(SubscribeResponse(initial={}))
# Re-sending flow tokens on the new stream
await write_called_queue.get()
await write_result_queue.put(None)
default_connection.write.assert_has_calls(
[
call(initial_request),
call(as_request(flow_1)),
call(initial_request),
call(
as_request(
FlowControlRequest(allowed_messages=115, allowed_bytes=115)
)
),
]
)
async def test_handle_reset_signal(
subscriber: Subscriber,
default_connection,
initial_request,
asyncio_sleep,
sleep_queues,
reset_handler,
):
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
flow = FlowControlRequest(allowed_messages=100, allowed_bytes=100)
message_1 = SequencedMessage(cursor=Cursor(offset=2), size_bytes=5)
message_2 = SequencedMessage(cursor=Cursor(offset=4), size_bytes=10)
# Ensure messages with earlier offsets can be handled post-reset.
message_3 = SequencedMessage(cursor=Cursor(offset=1), size_bytes=20)
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue
)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue
)
read_result_queue.put_nowait(SubscribeResponse(initial={}))
write_result_queue.put_nowait(None)
async with subscriber:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Send tokens.
flow_fut = asyncio.ensure_future(subscriber.allow_flow(flow))
assert not flow_fut.done()
# Handle the inline write since initial tokens are 100% of outstanding.
await write_called_queue.get()
await write_result_queue.put(None)
await flow_fut
default_connection.write.assert_has_calls(
[call(initial_request), call(as_request(flow))]
)
# Send messages to the subscriber.
await read_result_queue.put(as_response([message_1, message_2]))
# Read one message.
await read_called_queue.get()
assert (await subscriber.read()) == message_1
# Fail the connection with an error containing the RESET signal.
await read_called_queue.get()
await read_result_queue.put(make_reset_signal())
await sleep_queues[_MIN_BACKOFF_SECS].called.get()
await sleep_queues[_MIN_BACKOFF_SECS].results.put(None)
# Reinitialization.
await write_called_queue.get()
await write_result_queue.put(None)
await read_called_queue.get()
await read_result_queue.put(SubscribeResponse(initial={}))
# Re-sending flow tokens on the new stream.
await write_called_queue.get()
await write_result_queue.put(None)
reset_handler.handle_reset.assert_called_once()
default_connection.write.assert_has_calls(
[
call(initial_request),
call(as_request(flow)),
call(initial_request),
call(
as_request(
# Tokens for undelivered message_2 refilled.
FlowControlRequest(allowed_messages=99, allowed_bytes=95)
)
),
]
)
# Ensure the subscriber accepts an earlier message.
await read_result_queue.put(as_response([message_3]))
await read_called_queue.get()
assert (await subscriber.read()) == message_3
async def test_restart(
assigner: Assigner,
default_connection,
connection_factory,
initial_request,
asyncio_sleep,
sleep_queues,
):
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue)
write_result_queue.put_nowait(None)
async with assigner:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Wait for the first assignment
assign_fut1 = asyncio.ensure_future(assigner.get_assignment())
assert not assign_fut1.done()
partitions = {Partition(2), Partition(7)}
# Send the first assignment.
await read_result_queue.put(as_response(partitions=partitions))
await read_called_queue.get()
assert (await assign_fut1) == partitions
# Get the next assignment: should attempt to send an ack on the stream
assign_fut2 = asyncio.ensure_future(assigner.get_assignment())
await write_called_queue.get()
default_connection.write.assert_has_calls(
[call(initial_request), call(ack_request())])
# Set up the next connection
conn2 = MagicMock(spec=Connection)
conn2.__aenter__.return_value = conn2
connection_factory.new.return_value = conn2
write_called_queue_2 = asyncio.Queue()
write_result_queue_2 = asyncio.Queue()
conn2.write.side_effect = make_queue_waiter(write_called_queue_2,
write_result_queue_2)
read_called_queue_2 = asyncio.Queue()
read_result_queue_2 = asyncio.Queue()
conn2.read.side_effect = make_queue_waiter(read_called_queue_2,
read_result_queue_2)
# Fail the connection by failing the write call.
await write_result_queue.put(InternalServerError("failed"))
await sleep_queues[_MIN_BACKOFF_SECS].called.get()
await sleep_queues[_MIN_BACKOFF_SECS].results.put(None)
# Reinitialize
await write_called_queue_2.get()
write_result_queue_2.put_nowait(None)
conn2.write.assert_has_calls([call(initial_request)])
partitions = {Partition(5)}
# Send the second assignment on the new connection.
await read_called_queue_2.get()
await read_result_queue_2.put(as_response(partitions=partitions))
assert (await assign_fut2) == partitions
# No ack call ever made.
conn2.write.assert_has_calls([call(initial_request)])
async def test_message_receipt(
subscriber: Subscriber,
default_connection,
base_initial_subscribe,
initial_request,
asyncio_sleep,
sleep_queues,
):
write_called_queue = asyncio.Queue()
write_result_queue = asyncio.Queue()
flow = FlowControlRequest(allowed_messages=100, allowed_bytes=100)
message_1 = SequencedMessage(cursor=Cursor(offset=3), size_bytes=5)
message_2 = SequencedMessage(cursor=Cursor(offset=5), size_bytes=10)
default_connection.write.side_effect = make_queue_waiter(
write_called_queue, write_result_queue
)
read_called_queue = asyncio.Queue()
read_result_queue = asyncio.Queue()
default_connection.read.side_effect = make_queue_waiter(
read_called_queue, read_result_queue
)
read_result_queue.put_nowait(SubscribeResponse(initial={}))
write_result_queue.put_nowait(None)
async with subscriber:
# Set up connection
await write_called_queue.get()
await read_called_queue.get()
default_connection.write.assert_has_calls([call(initial_request)])
# Send tokens.
flow_fut = asyncio.ensure_future(subscriber.allow_flow(flow))
assert not flow_fut.done()
# Handle the inline write since initial tokens are 100% of outstanding.
await write_called_queue.get()
await write_result_queue.put(None)
await flow_fut
default_connection.write.assert_has_calls(
[call(initial_request), call(as_request(flow))]
)
message1_fut = asyncio.ensure_future(subscriber.read())
# Send messages to the subscriber.
await read_result_queue.put(as_response([message_1, message_2]))
# Wait for the next read call
await read_called_queue.get()
assert (await message1_fut) == message_1
assert (await subscriber.read()) == message_2
# Fail the connection with a retryable error
await read_called_queue.get()
await read_result_queue.put(InternalServerError("retryable"))
await sleep_queues[_MIN_BACKOFF_SECS].called.get()
await sleep_queues[_MIN_BACKOFF_SECS].results.put(None)
# Reinitialization
await write_called_queue.get()
seek_to_cursor_request = make_initial_subscribe_request(
base_initial_subscribe,
SeekRequest(cursor=Cursor(offset=message_2.cursor.offset + 1)),
)
default_connection.write.assert_has_calls(
[
call(initial_request),
call(as_request(flow)),
call(seek_to_cursor_request),
]
)
await write_result_queue.put(None)
await read_called_queue.get()
await read_result_queue.put(SubscribeResponse(initial={}))
# Re-sending flow tokens on the new stream.
await write_called_queue.get()
await write_result_queue.put(None)
default_connection.write.assert_has_calls(
[
call(initial_request),
call(as_request(flow)),
call(seek_to_cursor_request),
call(
as_request(
FlowControlRequest(allowed_messages=98, allowed_bytes=85)
)
),
]
)
