def test_recv_recover_race_condition(self): # This test checks the race condition where two threads recv() and # encounter an error and must re-open the stream. Only one thread # should succeed in doing so. error = ValueError() call_1 = CallStub([error, error]) call_2 = CallStub([1, 2]) start_rpc = mock.create_autospec(grpc.StreamStreamMultiCallable, instance=True, side_effect=[call_1, call_2]) recovered_event = threading.Event() def second_thread_main(): assert bidi_rpc.recv() == 2 second_thread = threading.Thread(target=second_thread_main) def should_recover(exception): assert exception == error if threading.current_thread() == second_thread: recovered_event.wait() return True bidi_rpc = bidi.ResumableBidiRpc(start_rpc, should_recover) bidi_rpc.open() second_thread.start() assert bidi_rpc.recv() == 1 recovered_event.set() assert bidi_rpc.call == call_2 assert bidi_rpc.is_active is True second_thread.join()
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 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_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_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_recv_failure(self): error = ValueError() call = CallStub([error]) start_rpc = mock.create_autospec(grpc.StreamStreamMultiCallable, instance=True, return_value=call) should_recover = mock.Mock(autospec=['__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 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(autospec=['__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_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(autospec=['__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_initial_state(self): bidi_rpc = bidi.ResumableBidiRpc(None, lambda _: True) assert bidi_rpc.is_active is False
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)