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()
