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