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)