@Override public void beforeStart(ClientCallStreamObserver<SimpleRequest> requestStream) { requestStream.setOnReadyHandler(() -> { if (numRequests.get() < TOTAL_NUM_MESSAGES) { numRequests.incrementAndGet(); requestStream.onNext(REQUEST); } }); } });
@Override public void onNext(SimpleResponse value) { if (numResponses.incrementAndGet() < CAPPED_NUM_MESSAGES) { requestStream.request(1); } else { if (!clientClosed.get()) { for (int i = 0; i < TOTAL_NUM_MESSAGES; i++) { requestStream.onNext(SimpleRequest.getDefaultInstance()); } requestStream.onCompleted(); clientClosed.set(true); } requestStream.request(1); } }
mRequestObserver.onNext(request);
/** * Checks send request is called on the request observer. */ @Test public void send() throws Exception { WriteRequest request = WriteRequest.newBuilder().build(); mStream.send(request, TIMEOUT); verify(mRequestObserver).onNext(request); }
/** * Validates the read request sent. * * @param client the worker client * @param offset the offset * @param length the length */ private void validateReadRequestSent(final BlockWorkerClient client, long offset, long length, boolean closed, int chunkSize) throws TimeoutException, InterruptedException { ArgumentCaptor<ReadRequest> requestCaptor = ArgumentCaptor.forClass(ReadRequest.class); verify(mRequestObserver, atLeastOnce()).onNext(requestCaptor.capture()); ArgumentCaptor<StreamObserver> captor = ArgumentCaptor.forClass(StreamObserver.class); verify(mClient).readBlock(captor.capture()); ReadRequest readRequest = requestCaptor.getValue(); captor.getValue().onCompleted(); verify(mRequestObserver, closed ? atLeastOnce() : never()).onCompleted(); assertTrue(readRequest != null); assertEquals(BLOCK_ID, readRequest.getBlockId()); assertEquals(offset, readRequest.getOffset()); assertEquals(length, readRequest.getLength()); assertEquals(chunkSize, readRequest.getChunkSize()); }
/** * Checks send after stream is ready. */ @Test public void sendAfterStreamReady() throws Exception { when(mRequestObserver.isReady()).thenReturn(false); doAnswer((args) -> { mOnReadyHandler = args.getArgumentAt(0, Runnable.class); return null; }).when(mRequestObserver).setOnReadyHandler(any(Runnable.class)); mResponseObserver.beforeStart(mRequestObserver); EXECUTOR.submit(() -> { try { // notify ready after a short period of time Thread.sleep(SHORT_TIMEOUT); when(mRequestObserver.isReady()).thenReturn(true); mOnReadyHandler.run(); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } }); WriteRequest request = WriteRequest.newBuilder().build(); mStream.send(request, TIMEOUT); verify(mRequestObserver).onNext(request); }
try { ArgumentCaptor<WriteRequest> requestCaptor = ArgumentCaptor.forClass(WriteRequest.class); verify(mRequestObserver, atLeastOnce()).onNext(requestCaptor.capture()); ArgumentCaptor<StreamObserver> captor = ArgumentCaptor.forClass(StreamObserver.class); verify(mClient).writeBlock(captor.capture());
if (req.isReady()) { numRequests.incrementAndGet(); req.onNext(REQUEST); } else { break;
try { ArgumentCaptor<WriteRequest> requestCaptor = ArgumentCaptor.forClass(WriteRequest.class); verify(mRequestObserver, atLeastOnce()).onNext(requestCaptor.capture()); ArgumentCaptor<StreamObserver> captor = ArgumentCaptor.forClass(StreamObserver.class); verify(mClient).writeBlock(captor.capture());