def test_on_response(): manager, _, dispatcher, _, scheduler = make_running_manager() manager._callback = mock.sentinel.callback # Set up the messages. response = types.StreamingPullResponse( received_messages=[ types.ReceivedMessage( ack_id='fack', message=types.PubsubMessage(data=b'foo', message_id='1') ), types.ReceivedMessage( ack_id='back', message=types.PubsubMessage(data=b'bar', message_id='2') ), ], ) # 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)
def test__on_response_no_leaser_overload(): manager, _, dispatcher, leaser, _, scheduler = make_running_manager() manager._callback = mock.sentinel.callback # Set up the messages. response = types.StreamingPullResponse(received_messages=[ types.ReceivedMessage(ack_id="fack", message=types.PubsubMessage(data=b"foo", message_id="1")), types.ReceivedMessage(ack_id="back", message=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_nacks_on_error(): # Create a callback that always errors. callback = mock.Mock(spec=(), side_effect=ValueError) executor = mock.create_autospec(futures.Executor, instance=True) executor.submit.side_effect = _callback_side_effect policy = create_and_open_policy(callback, executor=executor) # Set up the messages. message = types.PubsubMessage(data=b'foo', message_id='1') response = types.StreamingPullResponse(received_messages=[ types.ReceivedMessage(ack_id='fack', message=message), ], ) # Actually run the method and prove that nack is called because the # callback errored. policy.on_response(response) # Make sure the callback was executed. callback.assert_called_once_with(mock.ANY) # Process outstanding requests, the callback should've queued a nack # request. nack_patch = mock.patch.object(policy, 'nack', autospec=True) with nack_patch as nack: policy.dispatch_callback(policy._request_queue.queue) nack.assert_called_once_with( [base.NackRequest('fack', message.ByteSize())])
def test__on_response_delivery_attempt(): manager, _, dispatcher, leaser, _, scheduler = make_running_manager() manager._callback = mock.sentinel.callback # Set up the messages. response = types.StreamingPullResponse(received_messages=[ types.ReceivedMessage(ack_id="fack", message=types.PubsubMessage(data=b"foo", message_id="1")), types.ReceivedMessage( ack_id="back", message=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(): # Create mock Executor so we can verify calls to executor.submit(). executor = mock.create_autospec(futures.Executor, instance=True) callback = mock.Mock(spec=()) policy = create_and_open_policy(callback, executor=executor) # Set up the messages. response = types.StreamingPullResponse(received_messages=[ types.ReceivedMessage(ack_id='fack', message=types.PubsubMessage(data=b'foo', message_id='1')), types.ReceivedMessage(ack_id='back', message=types.PubsubMessage(data=b'bar', message_id='2')), ], ) # Actually run the method and prove that modack and executor.submit # are called in the expected way. modack_patch = mock.patch.object(policy, 'modify_ack_deadline', autospec=True) with modack_patch as modack: policy.on_response(response) modack.assert_called_once_with( [base.ModAckRequest('fack', 10), base.ModAckRequest('back', 10)]) submit_calls = [m for m in executor.method_calls if m[0] == 'submit'] assert len(submit_calls) == 2 for call in submit_calls: assert call[1][0] == policy._callback assert isinstance(call[1][1], message.Message)
def test_on_response(): callback = mock.Mock(spec=()) # Set up the policy. policy = create_policy() policy._callback = callback # Set up the messages to send. messages = ( types.PubsubMessage(data=b'foo', message_id='1'), types.PubsubMessage(data=b'bar', message_id='2'), ) # Set up a valid response. response = types.StreamingPullResponse(received_messages=[ { 'ack_id': 'fack', 'message': messages[0] }, { 'ack_id': 'back', 'message': messages[1] }, ], ) # Actually run the method and prove that the callback was # called in the expected way. policy.on_response(response) assert callback.call_count == 2 for call in callback.mock_calls: assert isinstance(call[1][0], message.Message)
def test__on_response_with_leaser_overload(): manager, _, dispatcher, leaser, _, scheduler = make_running_manager() manager._callback = mock.sentinel.callback # Set up the messages. response = types.StreamingPullResponse( received_messages=[ types.ReceivedMessage( ack_id="fack", message=types.PubsubMessage(data=b"foo", message_id="1") ), types.ReceivedMessage( ack_id="back", message=types.PubsubMessage(data=b"bar", message_id="2") ), types.ReceivedMessage( ack_id="zack", message=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(): callback = mock.Mock(spec=()) # Create mock ThreadPoolExecutor, pass into create_policy(), and verify # that both executor.submit() and future.add_done_callback are called # twice. future = mock.Mock() attrs = {'submit.return_value': future} executor = mock.Mock(**attrs) # Set up the policy. policy = create_policy(executor=executor) policy._callback = callback # Set up the messages to send. messages = ( types.PubsubMessage(data=b'foo', message_id='1'), types.PubsubMessage(data=b'bar', message_id='2'), ) # Set up a valid response. response = types.StreamingPullResponse(received_messages=[ { 'ack_id': 'fack', 'message': messages[0] }, { 'ack_id': 'back', 'message': messages[1] }, ], ) # Actually run the method and prove that modack and executor.submit # are called in the expected way. modack_patch = mock.patch.object(policy, 'modify_ack_deadline', autospec=True) with modack_patch as modack: policy.on_response(response) modack.assert_called_once_with( [base.ModAckRequest('fack', 10), base.ModAckRequest('back', 10)]) submit_calls = [m for m in executor.method_calls if m[0] == 'submit'] assert len(submit_calls) == 2 for call in submit_calls: assert call[1][0] == callback assert isinstance(call[1][1], message.Message)
def test_on_response(): callback = mock.Mock(spec=()) # Create mock ThreadPoolExecutor, pass into create_policy(), and verify # that both executor.submit() and future.add_done_callback are called # twice. future = mock.Mock() attrs = {'submit.return_value': future} executor = mock.Mock(**attrs) # Set up the policy. policy = create_policy(executor=executor) policy._callback = callback # Set up the messages to send. messages = ( types.PubsubMessage(data=b'foo', message_id='1'), types.PubsubMessage(data=b'bar', message_id='2'), ) # Set up a valid response. response = types.StreamingPullResponse(received_messages=[ { 'ack_id': 'fack', 'message': messages[0] }, { 'ack_id': 'back', 'message': messages[1] }, ], ) # Actually run the method and prove that executor.submit and # future.add_done_callback were called in the expected way. policy.on_response(response) submit_calls = [m for m in executor.method_calls if m[0] == 'submit'] assert len(submit_calls) == 2 for call in submit_calls: assert call[1][0] == callback assert isinstance(call[1][1], message.Message) add_done_callback_calls = [ m for m in future.method_calls if m[0] == 'add_done_callback' ] assert len(add_done_callback_calls) == 2 for call in add_done_callback_calls: assert call[1][0] == thread._callback_completed
def test__on_response_with_ordering_keys(): manager, _, dispatcher, leaser, _, scheduler = make_running_manager() manager._callback = mock.sentinel.callback # Set up the messages. response = types.StreamingPullResponse(received_messages=[ types.ReceivedMessage( ack_id="fack", message=types.PubsubMessage( data=b"foo", message_id="1", ordering_key=""), ), types.ReceivedMessage( ack_id="back", message=types.PubsubMessage( data=b"bar", message_id="2", ordering_key="key1"), ), types.ReceivedMessage( ack_id="zack", message=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