def test_send_unary_empty():
manager = make_manager()
manager.send(gapic_types.StreamingPullRequest())
manager._client.acknowledge.assert_not_called()
manager._client.modify_ack_deadline.assert_not_called()
def test_send_unary():
manager = make_manager()
manager._UNARY_REQUESTS = True
manager.send(
gapic_types.StreamingPullRequest(
ack_ids=["ack_id1", "ack_id2"],
modify_deadline_ack_ids=["ack_id3", "ack_id4", "ack_id5"],
modify_deadline_seconds=[10, 20, 20],
))
manager._client.acknowledge.assert_called_once_with(
subscription=manager._subscription, ack_ids=["ack_id1", "ack_id2"])
manager._client.modify_ack_deadline.assert_has_calls(
[
mock.call(
subscription=manager._subscription,
ack_ids=["ack_id3"],
ack_deadline_seconds=10,
),
mock.call(
subscription=manager._subscription,
ack_ids=["ack_id4", "ack_id5"],
ack_deadline_seconds=20,
),
],
any_order=True,
)
def heartbeat(self):
"""Sends an empty request over the streaming pull RPC.
This always sends over the stream, regardless of if
``self._UNARY_REQUESTS`` is set or not.
"""
if self._rpc is not None and self._rpc.is_active:
self._rpc.send(gapic_types.StreamingPullRequest())
def test_heartbeat():
manager = make_manager()
manager._rpc = mock.create_autospec(bidi.BidiRpc, instance=True)
manager._rpc.is_active = True
manager.heartbeat()
manager._rpc.send.assert_called_once_with(gapic_types.StreamingPullRequest())
def heartbeat(self):
"""Sends an empty request over the streaming pull RPC.
Returns:
bool: If a heartbeat request has actually been sent.
"""
if self._rpc is not None and self._rpc.is_active:
self._rpc.send(gapic_types.StreamingPullRequest())
return True
return False
def test_send_unary_api_call_error(caplog):
caplog.set_level(logging.DEBUG)
manager = make_manager()
error = exceptions.GoogleAPICallError("The front fell off")
manager._client.acknowledge.side_effect = error
manager.send(gapic_types.StreamingPullRequest(ack_ids=["ack_id1", "ack_id2"]))
assert "The front fell off" in caplog.text
def test_modify_ack_deadline():
manager = mock.create_autospec(streaming_pull_manager.StreamingPullManager,
instance=True)
dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)
items = [requests.ModAckRequest(ack_id="ack_id_string", seconds=60)]
dispatcher_.modify_ack_deadline(items)
manager.send.assert_called_once_with(
gapic_types.StreamingPullRequest(
modify_deadline_ack_ids=["ack_id_string"],
modify_deadline_seconds=[60]))
def test_send_unary_retry_error(caplog):
caplog.set_level(logging.DEBUG)
manager, _, _, _, _, _ = make_running_manager()
error = exceptions.RetryError("Too long a transient error",
cause=Exception("Out of time!"))
manager._client.acknowledge.side_effect = error
with pytest.raises(exceptions.RetryError):
manager.send(
gapic_types.StreamingPullRequest(ack_ids=["ack_id1", "ack_id2"]))
assert "RetryError while sending unary RPC" in caplog.text
assert "signaled streaming pull manager shutdown" in caplog.text
def test_nack():
manager = mock.create_autospec(streaming_pull_manager.StreamingPullManager,
instance=True)
dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)
items = [
requests.NackRequest(ack_id="ack_id_string",
byte_size=10,
ordering_key="")
]
dispatcher_.nack(items)
manager.send.assert_called_once_with(
gapic_types.StreamingPullRequest(
modify_deadline_ack_ids=["ack_id_string"],
modify_deadline_seconds=[0]))
def test_ack_no_time():
manager = mock.create_autospec(streaming_pull_manager.StreamingPullManager,
instance=True)
dispatcher_ = dispatcher.Dispatcher(manager, mock.sentinel.queue)
items = [
requests.AckRequest(ack_id="ack_id_string",
byte_size=0,
time_to_ack=None,
ordering_key="")
]
dispatcher_.ack(items)
manager.send.assert_called_once_with(
gapic_types.StreamingPullRequest(ack_ids=["ack_id_string"]))
manager.ack_histogram.add.assert_not_called()
def modify_ack_deadline(self, items):
"""Modify the ack deadline for the given messages.
Args:
items(Sequence[ModAckRequest]): The items to modify.
"""
# We must potentially split the request into multiple smaller requests
# to avoid the server-side max request size limit.
ack_ids = (item.ack_id for item in items)
seconds = (item.seconds for item in items)
total_chunks = int(math.ceil(len(items) / _ACK_IDS_BATCH_SIZE))
for _ in range(total_chunks):
request = gapic_types.StreamingPullRequest(
modify_deadline_ack_ids=itertools.islice(
ack_ids, _ACK_IDS_BATCH_SIZE),
modify_deadline_seconds=itertools.islice(
seconds, _ACK_IDS_BATCH_SIZE),
)
self._manager.send(request)
def _get_initial_request(self, stream_ack_deadline_seconds):
"""Return the initial request for the RPC.
This defines the initial request that must always be sent to Pub/Sub
immediately upon opening the subscription.
Args:
stream_ack_deadline_seconds (int):
The default message acknowledge deadline for the stream.
Returns:
google.pubsub_v1.types.StreamingPullRequest: A request
suitable for being the first request on the stream (and not
suitable for any other purpose).
"""
# Any ack IDs that are under lease management need to have their
# deadline extended immediately.
if self._leaser is not None:
# Explicitly copy the list, as it could be modified by another
# thread.
lease_ids = list(self._leaser.ack_ids)
else:
lease_ids = []
# Put the request together.
request = gapic_types.StreamingPullRequest(
modify_deadline_ack_ids=list(lease_ids),
modify_deadline_seconds=[self.ack_deadline] * len(lease_ids),
stream_ack_deadline_seconds=stream_ack_deadline_seconds,
subscription=self._subscription,
client_id=self._client_id,
max_outstanding_messages=(
0 if self._use_legacy_flow_control else self._flow_control.max_messages
),
max_outstanding_bytes=(
0 if self._use_legacy_flow_control else self._flow_control.max_bytes
),
)
# Return the initial request.
return request
def ack(self, items):
"""Acknowledge the given messages.
Args:
items(Sequence[AckRequest]): The items to acknowledge.
"""
# If we got timing information, add it to the histogram.
for item in items:
time_to_ack = item.time_to_ack
if time_to_ack is not None:
self._manager.ack_histogram.add(time_to_ack)
# We must potentially split the request into multiple smaller requests
# to avoid the server-side max request size limit.
ack_ids = (item.ack_id for item in items)
total_chunks = int(math.ceil(len(items) / _ACK_IDS_BATCH_SIZE))
for _ in range(total_chunks):
request = gapic_types.StreamingPullRequest(
ack_ids=itertools.islice(ack_ids, _ACK_IDS_BATCH_SIZE))
self._manager.send(request)
# Remove the message from lease management.
self.drop(items)