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) ) ), ] )