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
