def test_reopen_failure_on_rpc_restart(self): error1 = ValueError('1') error2 = ValueError('2') call = CallStub([error1]) # Invoking start RPC a second time will trigger an error. start_rpc = mock.create_autospec( grpc.StreamStreamMultiCallable, instance=True, side_effect=[call, error2]) should_recover = mock.Mock(spec=['__call__'], return_value=True) callback = mock.Mock(spec=['__call__']) bidi_rpc = bidi.ResumableBidiRpc(start_rpc, should_recover) bidi_rpc.add_done_callback(callback) bidi_rpc.open() with pytest.raises(ValueError) as exc_info: bidi_rpc.recv() assert exc_info.value == error2 should_recover.assert_called_once_with(error1) assert bidi_rpc.call is None assert bidi_rpc.is_active is False callback.assert_called_once_with(error2)
def test_send_terminate(self): cancellation = ValueError() call_1 = CallStub([cancellation], active=False) call_2 = CallStub([]) start_rpc = mock.create_autospec(grpc.StreamStreamMultiCallable, instance=True, side_effect=[call_1, call_2]) should_recover = mock.Mock(spec=["__call__"], return_value=False) should_terminate = mock.Mock(spec=["__call__"], return_value=True) bidi_rpc = bidi.ResumableBidiRpc(start_rpc, should_recover, should_terminate=should_terminate) bidi_rpc.open() bidi_rpc.send(mock.sentinel.request) assert bidi_rpc.pending_requests == 1 assert bidi_rpc._request_queue.get() is None should_recover.assert_not_called() should_terminate.assert_called_once_with(cancellation) assert bidi_rpc.call == call_1 assert bidi_rpc.is_active is False assert call_1.cancelled is True
def open(self, callback, on_callback_error): """Begin consuming messages. Args: callback (Callable[None, google.cloud.pubsub_v1.message.Message]): A callback that will be called for each message received on the stream. on_callback_error (Callable[Exception]): A callable that will be called if an exception is raised in the provided `callback`. """ if self.is_active: raise ValueError("This manager is already open.") if self._closed: raise ValueError("This manager has been closed and can not be re-used.") self._callback = functools.partial( _wrap_callback_errors, callback, on_callback_error ) # Create the RPC stream_ack_deadline_seconds = self.ack_histogram.percentile(99) get_initial_request = functools.partial( self._get_initial_request, stream_ack_deadline_seconds ) self._rpc = bidi.ResumableBidiRpc( start_rpc=self._client.api.streaming_pull, initial_request=get_initial_request, should_recover=self._should_recover, should_terminate=self._should_terminate, throttle_reopen=True, ) self._rpc.add_done_callback(self._on_rpc_done) _LOGGER.debug( "Creating a stream, default ACK deadline set to {} seconds.".format( stream_ack_deadline_seconds ) ) # Create references to threads self._dispatcher = dispatcher.Dispatcher(self, self._scheduler.queue) self._consumer = bidi.BackgroundConsumer(self._rpc, self._on_response) self._leaser = leaser.Leaser(self) self._heartbeater = heartbeater.Heartbeater(self) # Start the thread to pass the requests. self._dispatcher.start() # Start consuming messages. self._consumer.start() # Start the lease maintainer thread. self._leaser.start() # Start the stream heartbeater thread. self._heartbeater.start()
def test_done_callbacks_non_recoverable(self): bidi_rpc = bidi.ResumableBidiRpc(None, lambda _: False) callback = mock.Mock(spec=['__call__']) bidi_rpc.add_done_callback(callback) bidi_rpc._on_call_done(mock.sentinel.future) callback.assert_called_once_with(mock.sentinel.future)
def test_done_callbacks_recoverable(self): start_rpc = mock.create_autospec( grpc.StreamStreamMultiCallable, instance=True) bidi_rpc = bidi.ResumableBidiRpc(start_rpc, lambda _: True) callback = mock.Mock(spec=['__call__']) bidi_rpc.add_done_callback(callback) bidi_rpc._on_call_done(mock.sentinel.future) callback.assert_not_called() start_rpc.assert_called_once() assert bidi_rpc.is_active
def test_done_callbacks_non_recoverable(self): start_rpc = mock.create_autospec(grpc.StreamStreamMultiCallable, instance=True) should_recover = mock.Mock(spec=["__call__"], return_value=False) bidi_rpc = bidi.ResumableBidiRpc(start_rpc, should_recover) callback = mock.Mock(spec=["__call__"]) bidi_rpc.add_done_callback(callback) bidi_rpc._on_call_done(mock.sentinel.future) callback.assert_called_once_with(mock.sentinel.future) should_recover.assert_called_once_with(mock.sentinel.future) assert not bidi_rpc.is_active
def test_ctor_defaults(self): start_rpc = mock.Mock() should_recover = mock.Mock() bidi_rpc = bidi.ResumableBidiRpc(start_rpc, should_recover) assert bidi_rpc.is_active is False assert bidi_rpc._finalized is False assert bidi_rpc._start_rpc is start_rpc assert bidi_rpc._should_recover is should_recover assert bidi_rpc._should_terminate is bidi._never_terminate assert bidi_rpc._initial_request is None assert bidi_rpc._rpc_metadata is None assert bidi_rpc._reopen_throttle is None
def test_recv_recover_already_recovered(self): call_1 = CallStub([]) call_2 = CallStub([]) start_rpc = mock.create_autospec(grpc.StreamStreamMultiCallable, instance=True, side_effect=[call_1, call_2]) bidi_rpc = bidi.ResumableBidiRpc(start_rpc, lambda _: True) bidi_rpc.open() bidi_rpc._reopen() assert bidi_rpc.call is call_1 assert bidi_rpc.is_active is True
def test_finalize_idempotent(self): error1 = ValueError("1") error2 = ValueError("2") callback = mock.Mock(spec=["__call__"]) should_recover = mock.Mock(spec=["__call__"], return_value=False) bidi_rpc = bidi.ResumableBidiRpc(mock.sentinel.start_rpc, should_recover) bidi_rpc.add_done_callback(callback) bidi_rpc._on_call_done(error1) bidi_rpc._on_call_done(error2) callback.assert_called_once_with(error1)
def test_using_throttle_on_reopen_requests(self): call = CallStub([]) start_rpc = mock.create_autospec(grpc.StreamStreamMultiCallable, instance=True, return_value=call) should_recover = mock.Mock(spec=["__call__"], return_value=True) bidi_rpc = bidi.ResumableBidiRpc(start_rpc, should_recover, throttle_reopen=True) patcher = mock.patch.object(bidi_rpc._reopen_throttle.__class__, "__enter__") with patcher as mock_enter: bidi_rpc._reopen() mock_enter.assert_called_once()
def test_done_callbacks_terminate(self): cancellation = mock.Mock() start_rpc = mock.Mock() should_recover = mock.Mock(spec=["__call__"], return_value=True) should_terminate = mock.Mock(spec=["__call__"], return_value=True) bidi_rpc = bidi.ResumableBidiRpc(start_rpc, should_recover, should_terminate=should_terminate) callback = mock.Mock(spec=["__call__"]) bidi_rpc.add_done_callback(callback) bidi_rpc._on_call_done(cancellation) should_terminate.assert_called_once_with(cancellation) should_recover.assert_not_called() callback.assert_called_once_with(cancellation) assert not bidi_rpc.is_active
def test_recv_failure(self): error = ValueError() call = CallStub([error]) start_rpc = mock.create_autospec(grpc.StreamStreamMultiCallable, instance=True, return_value=call) should_recover = mock.Mock(spec=["__call__"], return_value=False) bidi_rpc = bidi.ResumableBidiRpc(start_rpc, should_recover) bidi_rpc.open() with pytest.raises(ValueError) as exc_info: bidi_rpc.recv() assert exc_info.value == error should_recover.assert_called_once_with(error) assert bidi_rpc.call == call assert bidi_rpc.is_active is False assert call.cancelled is True
def open(self, callback): """Begin consuming messages. Args: callback (Callable[None, google.cloud.pubsub_v1.message.Messages]): A callback that will be called for each message received on the stream. """ if self.is_active: raise ValueError('This manager is already open.') if self._closed: raise ValueError( 'This manager has been closed and can not be re-used.') self._callback = functools.partial(_wrap_callback_errors, callback) # Create the RPC self._rpc = bidi.ResumableBidiRpc( start_rpc=self._client.api.streaming_pull, initial_request=self._get_initial_request, should_recover=self._should_recover) self._rpc.add_done_callback(self._on_rpc_done) # Create references to threads self._dispatcher = dispatcher.Dispatcher(self, self._scheduler.queue) self._consumer = bidi.BackgroundConsumer( self._rpc, self._on_response) self._leaser = leaser.Leaser(self) self._heartbeater = heartbeater.Heartbeater(self) # Start the thread to pass the requests. self._dispatcher.start() # Start consuming messages. self._consumer.start() # Start the lease maintainer thread. self._leaser.start() # Start the stream heartbeater thread. self._heartbeater.start()
def test_send_recover(self): error = ValueError() call_1 = CallStub([error], active=False) call_2 = CallStub([]) start_rpc = mock.create_autospec(grpc.StreamStreamMultiCallable, instance=True, side_effect=[call_1, call_2]) should_recover = mock.Mock(spec=["__call__"], return_value=True) bidi_rpc = bidi.ResumableBidiRpc(start_rpc, should_recover) bidi_rpc.open() bidi_rpc.send(mock.sentinel.request) assert bidi_rpc.pending_requests == 1 assert bidi_rpc._request_queue.get() is mock.sentinel.request should_recover.assert_called_once_with(error) assert bidi_rpc.call == call_2 assert bidi_rpc.is_active is True
def test_recv_recover(self): error = ValueError() call_1 = CallStub([1, error]) call_2 = CallStub([2, 3]) start_rpc = mock.create_autospec(grpc.StreamStreamMultiCallable, instance=True, side_effect=[call_1, call_2]) should_recover = mock.Mock(spec=["__call__"], return_value=True) bidi_rpc = bidi.ResumableBidiRpc(start_rpc, should_recover) bidi_rpc.open() values = [] for n in range(3): values.append(bidi_rpc.recv()) assert values == [1, 2, 3] should_recover.assert_called_once_with(error) assert bidi_rpc.call == call_2 assert bidi_rpc.is_active is True
def test_send_failure(self): error = ValueError() call = CallStub([error], active=False) start_rpc = mock.create_autospec(grpc.StreamStreamMultiCallable, instance=True, return_value=call) should_recover = mock.Mock(spec=['__call__'], return_value=False) bidi_rpc = bidi.ResumableBidiRpc(start_rpc, should_recover) bidi_rpc.open() with pytest.raises(ValueError) as exc_info: bidi_rpc.send(mock.sentinel.request) assert exc_info.value == error should_recover.assert_called_once_with(error) assert bidi_rpc.call == call assert bidi_rpc.is_active is False assert call.cancelled is True assert bidi_rpc.pending_requests == 1 assert bidi_rpc._request_queue.get() is None
def test_recv_terminate(self): cancellation = ValueError() call = CallStub([cancellation]) start_rpc = mock.create_autospec(grpc.StreamStreamMultiCallable, instance=True, return_value=call) should_recover = mock.Mock(spec=["__call__"], return_value=False) should_terminate = mock.Mock(spec=["__call__"], return_value=True) bidi_rpc = bidi.ResumableBidiRpc(start_rpc, should_recover, should_terminate=should_terminate) bidi_rpc.open() bidi_rpc.recv() should_recover.assert_not_called() should_terminate.assert_called_once_with(cancellation) assert bidi_rpc.call == call assert bidi_rpc.is_active is False assert call.cancelled is True
def test_ctor_explicit(self): start_rpc = mock.Mock() should_recover = mock.Mock() should_terminate = mock.Mock() initial_request = mock.Mock() metadata = {"x-foo": "bar"} bidi_rpc = bidi.ResumableBidiRpc( start_rpc, should_recover, should_terminate=should_terminate, initial_request=initial_request, metadata=metadata, throttle_reopen=True, ) assert bidi_rpc.is_active is False assert bidi_rpc._finalized is False assert bidi_rpc._should_recover is should_recover assert bidi_rpc._should_terminate is should_terminate assert bidi_rpc._initial_request is initial_request assert bidi_rpc._rpc_metadata == metadata assert isinstance(bidi_rpc._reopen_throttle, bidi._Throttle)
def test_close(self): call = mock.create_autospec(_CallAndFuture, instance=True) def cancel_side_effect(): call.is_active.return_value = False call.cancel.side_effect = cancel_side_effect start_rpc = mock.create_autospec(grpc.StreamStreamMultiCallable, instance=True, return_value=call) should_recover = mock.Mock(spec=["__call__"], return_value=False) bidi_rpc = bidi.ResumableBidiRpc(start_rpc, should_recover) bidi_rpc.open() bidi_rpc.close() should_recover.assert_not_called() call.cancel.assert_called_once() assert bidi_rpc.call == call assert bidi_rpc.is_active is False # ensure the request queue was signaled to stop. assert bidi_rpc.pending_requests == 1 assert bidi_rpc._request_queue.get() is None assert bidi_rpc._finalized
def test_recv_not_open(self): bidi_rpc = bidi.ResumableBidiRpc(None, lambda _: False) with pytest.raises(ValueError): bidi_rpc.recv()
def test_send_not_open(self): bidi_rpc = bidi.ResumableBidiRpc(None, lambda _: False) with pytest.raises(ValueError): bidi_rpc.send(mock.sentinel.request)
def test_initial_state(self): callback = mock.Mock() callback.return_value = True bidi_rpc = bidi.ResumableBidiRpc(None, callback) assert bidi_rpc.is_active is False
def test_initial_state(self): bidi_rpc = bidi.ResumableBidiRpc(None, lambda _: True) assert bidi_rpc.is_active is False
def open(self, callback, on_callback_error): """Begin consuming messages. Args: callback (Callable[None, google.cloud.pubsub_v1.message.Message]): A callback that will be called for each message received on the stream. on_callback_error (Callable[Exception]): A callable that will be called if an exception is raised in the provided `callback`. """ if self.is_active: raise ValueError("This manager is already open.") if self._closed: raise ValueError( "This manager has been closed and can not be re-used.") self._callback = functools.partial(_wrap_callback_errors, callback, on_callback_error) # Create the RPC # We must use a fixed value for the ACK deadline, as we cannot read it # from the subscription. The latter would require `pubsub.subscriptions.get` # permission, which is not granted to the default subscriber role # `roles/pubsub.subscriber`. # See also https://github.com/googleapis/google-cloud-python/issues/9339 # # When dynamic lease management is enabled for the "on hold" messages, # the default stream ACK deadline should again be set based on the # historic ACK timing data, i.e. `self.ack_histogram.percentile(99)`. stream_ack_deadline_seconds = _DEFAULT_STREAM_ACK_DEADLINE get_initial_request = functools.partial(self._get_initial_request, stream_ack_deadline_seconds) self._rpc = bidi.ResumableBidiRpc( start_rpc=self._client.api.streaming_pull, initial_request=get_initial_request, should_recover=self._should_recover, should_terminate=self._should_terminate, throttle_reopen=True, ) self._rpc.add_done_callback(self._on_rpc_done) _LOGGER.debug( "Creating a stream, default ACK deadline set to {} seconds.". format(stream_ack_deadline_seconds)) # Create references to threads self._dispatcher = dispatcher.Dispatcher(self, self._scheduler.queue) self._consumer = bidi.BackgroundConsumer(self._rpc, self._on_response) self._leaser = leaser.Leaser(self) self._heartbeater = heartbeater.Heartbeater(self) # Start the thread to pass the requests. self._dispatcher.start() # Start consuming messages. self._consumer.start() # Start the lease maintainer thread. self._leaser.start() # Start the stream heartbeater thread. self._heartbeater.start()