def test_wake_on_error(self):
should_continue = threading.Event()
bidi_rpc = mock.create_autospec(bidi.BidiRpc, instance=True)
bidi_rpc.is_active = True
bidi_rpc.add_done_callback.side_effect = (
lambda _: should_continue.set())
consumer = bidi.BackgroundConsumer(bidi_rpc, mock.sentinel.on_response)
# Start the consumer paused, which should immediately put it into wait
# state.
consumer.pause()
consumer.start()
# Wait for add_done_callback to be called
should_continue.wait()
bidi_rpc.add_done_callback.assert_called_once_with(
consumer._on_call_done)
# The consumer should now be blocked on waiting to be unpaused.
assert consumer.is_active
assert consumer.is_paused
# Trigger the done callback, it should unpause the consumer and cause
# it to exit.
bidi_rpc.is_active = False
consumer._on_call_done(bidi_rpc)
# It may take a few cycles for the thread to exit.
while consumer.is_active:
pass
def test_pause_resume_and_close(self):
# This test is relatively complex. It attempts to start the consumer,
# consume one item, pause the consumer, check the state of the world,
# then resume the consumer. Doing this in a deterministic fashion
# requires a bit more mocking and patching than usual.
bidi_rpc = mock.create_autospec(bidi.BidiRpc, instance=True)
bidi_rpc.is_active = True
def close_side_effect():
bidi_rpc.is_active = False
bidi_rpc.close.side_effect = close_side_effect
# These are used to coordinate the two threads to ensure deterministic
# execution.
should_continue = threading.Event()
responses_and_events = {
mock.sentinel.response_1: threading.Event(),
mock.sentinel.response_2: threading.Event()
}
bidi_rpc.recv.side_effect = [
mock.sentinel.response_1, mock.sentinel.response_2]
recved_responses = []
consumer = None
def on_response(response):
if response == mock.sentinel.response_1:
consumer.pause()
recved_responses.append(response)
responses_and_events[response].set()
should_continue.wait()
consumer = bidi.BackgroundConsumer(bidi_rpc, on_response)
consumer.start()
# Wait for the first response to be recved.
responses_and_events[mock.sentinel.response_1].wait()
# Ensure only one item has been recved and that the consumer is paused.
assert recved_responses == [mock.sentinel.response_1]
assert consumer.is_paused is True
assert consumer.is_active is True
# Unpause the consumer, wait for the second item, then close the
# consumer.
should_continue.set()
consumer.resume()
responses_and_events[mock.sentinel.response_2].wait()
assert recved_responses == [
mock.sentinel.response_1, mock.sentinel.response_2]
consumer.stop()
assert consumer.is_active is False
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)
# Start the thread to pass the requests.
self._dispatcher = dispatcher.Dispatcher(self, self._scheduler.queue)
self._dispatcher.start()
# Start consuming messages.
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)
self._consumer = bidi.BackgroundConsumer(
self._rpc, self._on_response)
self._consumer.start()
# Start the lease maintainer thread.
self._leaser = leaser.Leaser(self)
self._leaser.start()
def test_consumer_unexpected_error(self, caplog):
caplog.set_level(logging.DEBUG)
bidi_rpc = mock.create_autospec(bidi.BidiRpc, instance=True)
bidi_rpc.is_active = True
bidi_rpc.recv.side_effect = ValueError()
on_response = mock.Mock(spec=['__call__'])
consumer = bidi.BackgroundConsumer(bidi_rpc, on_response)
consumer.start()
# Wait for the consumer's thread to exit.
while consumer.is_active:
pass
on_response.assert_not_called()
bidi_rpc.recv.assert_called_once()
assert 'caught unexpected exception' in caplog.text
def test_double_stop(self, caplog):
caplog.set_level(logging.DEBUG)
bidi_rpc = mock.create_autospec(bidi.BidiRpc, instance=True)
bidi_rpc.is_active = True
on_response = mock.Mock(spec=['__call__'])
def close_side_effect():
bidi_rpc.is_active = False
bidi_rpc.close.side_effect = close_side_effect
consumer = bidi.BackgroundConsumer(bidi_rpc, on_response)
consumer.start()
assert consumer.is_active is True
consumer.stop()
assert consumer.is_active is False
# calling stop twice should not result in an error.
consumer.stop()
def test_consume_once_then_exit(self):
bidi_rpc = mock.create_autospec(bidi.BidiRpc, instance=True)
bidi_rpc.is_active = True
bidi_rpc.recv.side_effect = [mock.sentinel.response_1]
recved = threading.Event()
def on_response(response):
assert response == mock.sentinel.response_1
bidi_rpc.is_active = False
recved.set()
consumer = bidi.BackgroundConsumer(bidi_rpc, on_response)
consumer.start()
recved.wait()
bidi_rpc.recv.assert_called_once()
assert bidi_rpc.is_active is False
consumer.stop()
bidi_rpc.close.assert_called_once()
assert consumer.is_active is False
