def test__on_response_delivery_attempt(): manager, _, dispatcher, leaser, _, scheduler = make_running_manager() manager._callback = mock.sentinel.callback # Set up the messages. response = gapic_types.StreamingPullResponse(received_messages=[ gapic_types.ReceivedMessage( ack_id="fack", message=gapic_types.PubsubMessage(data=b"foo", message_id="1"), ), gapic_types.ReceivedMessage( ack_id="back", message=gapic_types.PubsubMessage(data=b"bar", message_id="2"), delivery_attempt=6, ), ]) # adjust message bookkeeping in leaser fake_leaser_add(leaser, init_msg_count=0, assumed_msg_size=42) manager._on_response(response) schedule_calls = scheduler.schedule.mock_calls assert len(schedule_calls) == 2 msg1 = schedule_calls[0][1][1] assert msg1.delivery_attempt is None msg2 = schedule_calls[1][1][1] assert msg2.delivery_attempt == 6
def test__on_response_modifies_ack_deadline(): manager, _, dispatcher, leaser, _, scheduler = make_running_manager() manager._callback = mock.sentinel.callback # Set up the messages. response = gapic_types.StreamingPullResponse(received_messages=[ gapic_types.ReceivedMessage( ack_id="ack_1", message=gapic_types.PubsubMessage(data=b"foo", message_id="1"), ), gapic_types.ReceivedMessage( ack_id="ack_2", message=gapic_types.PubsubMessage(data=b"bar", message_id="2"), ), ]) # adjust message bookkeeping in leaser fake_leaser_add(leaser, init_msg_count=0, assumed_msg_size=80) # Actually run the method and chack that correct MODACK value is used. with mock.patch.object(type(manager), "ack_deadline", new=mock.PropertyMock(return_value=18)): manager._on_response(response) dispatcher.modify_ack_deadline.assert_called_once_with([ requests.ModAckRequest("ack_1", 18), requests.ModAckRequest("ack_2", 18) ])
def test__on_response_with_leaser_overload(): manager, _, dispatcher, leaser, _, scheduler = make_running_manager() manager._callback = mock.sentinel.callback # Set up the messages. response = gapic_types.StreamingPullResponse( received_messages=[ gapic_types.ReceivedMessage( ack_id="fack", message=gapic_types.PubsubMessage(data=b"foo", message_id="1"), ), gapic_types.ReceivedMessage( ack_id="back", message=gapic_types.PubsubMessage(data=b"bar", message_id="2"), ), gapic_types.ReceivedMessage( ack_id="zack", message=gapic_types.PubsubMessage(data=b"baz", message_id="3"), ), ] ) # Adjust message bookkeeping in leaser. Pick 999 messages, which is just below # the default FlowControl.max_messages limit. fake_leaser_add(leaser, init_msg_count=999, assumed_msg_size=10) # Actually run the method and prove that modack and schedule # are called in the expected way. manager._on_response(response) # all messages should be added to the lease management and have their ACK # deadline extended, even those not dispatched to callbacks dispatcher.modify_ack_deadline.assert_called_once_with( [ requests.ModAckRequest("fack", 10), requests.ModAckRequest("back", 10), requests.ModAckRequest("zack", 10), ] ) # one message should be scheduled, the flow control limits allow for it schedule_calls = scheduler.schedule.mock_calls assert len(schedule_calls) == 1 call_args = schedule_calls[0][1] assert call_args[0] == mock.sentinel.callback assert isinstance(call_args[1], message.Message) assert call_args[1].message_id == "1" # the rest of the messages should have been put on hold assert manager._messages_on_hold.size == 2 while True: msg = manager._messages_on_hold.get() if msg is None: break else: assert isinstance(msg, message.Message) assert msg.message_id in ("2", "3")
def test__on_response_no_leaser_overload(): manager, _, dispatcher, leaser, _, scheduler = make_running_manager() manager._callback = mock.sentinel.callback # Set up the messages. response = gapic_types.StreamingPullResponse( received_messages=[ gapic_types.ReceivedMessage( ack_id="fack", message=gapic_types.PubsubMessage(data=b"foo", message_id="1"), ), gapic_types.ReceivedMessage( ack_id="back", message=gapic_types.PubsubMessage(data=b"bar", message_id="2"), ), ] ) # adjust message bookkeeping in leaser fake_leaser_add(leaser, init_msg_count=0, assumed_msg_size=42) # Actually run the method and prove that modack and schedule # are called in the expected way. manager._on_response(response) dispatcher.modify_ack_deadline.assert_called_once_with( [requests.ModAckRequest("fack", 10), requests.ModAckRequest("back", 10)] ) schedule_calls = scheduler.schedule.mock_calls assert len(schedule_calls) == 2 for call in schedule_calls: assert call[1][0] == mock.sentinel.callback assert isinstance(call[1][1], message.Message) # the leaser load limit not hit, no messages had to be put on hold assert manager._messages_on_hold.size == 0
def test__on_response_with_ordering_keys(): manager, _, dispatcher, leaser, _, scheduler = make_running_manager() manager._callback = mock.sentinel.callback # Set up the messages. response = gapic_types.StreamingPullResponse(received_messages=[ gapic_types.ReceivedMessage( ack_id="fack", message=gapic_types.PubsubMessage( data=b"foo", message_id="1", ordering_key=""), ), gapic_types.ReceivedMessage( ack_id="back", message=gapic_types.PubsubMessage( data=b"bar", message_id="2", ordering_key="key1"), ), gapic_types.ReceivedMessage( ack_id="zack", message=gapic_types.PubsubMessage( data=b"baz", message_id="3", ordering_key="key1"), ), ]) # Make leaser with zero initial messages, so we don't test lease management # behavior. fake_leaser_add(leaser, init_msg_count=0, assumed_msg_size=10) # Actually run the method and prove that modack and schedule are called in # the expected way. manager._on_response(response) # All messages should be added to the lease management and have their ACK # deadline extended, even those not dispatched to callbacks. dispatcher.modify_ack_deadline.assert_called_once_with([ requests.ModAckRequest("fack", 10), requests.ModAckRequest("back", 10), requests.ModAckRequest("zack", 10), ]) # The first two messages should be scheduled, The third should be put on # hold because it's blocked by the completion of the second, which has the # same ordering key. schedule_calls = scheduler.schedule.mock_calls assert len(schedule_calls) == 2 call_args = schedule_calls[0][1] assert call_args[0] == mock.sentinel.callback assert isinstance(call_args[1], message.Message) assert call_args[1].message_id == "1" call_args = schedule_calls[1][1] assert call_args[0] == mock.sentinel.callback assert isinstance(call_args[1], message.Message) assert call_args[1].message_id == "2" # Message 3 should have been put on hold. assert manager._messages_on_hold.size == 1 # No messages available because message 2 (with "key1") has not completed yet. assert manager._messages_on_hold.get() is None # Complete message 2 (with "key1"). manager.activate_ordering_keys(["key1"]) # Completing message 2 should release message 3. schedule_calls = scheduler.schedule.mock_calls assert len(schedule_calls) == 3 call_args = schedule_calls[2][1] assert call_args[0] == mock.sentinel.callback assert isinstance(call_args[1], message.Message) assert call_args[1].message_id == "3" # No messages available in the queue. assert manager._messages_on_hold.get() is None