GrpcWindmillServer extends WindmillServerStub { @Override public GetWorkStream getWorkStream(GetWorkRequest request, WorkItemReceiver receiver) { return new GrpcGetWorkStream( GetWorkRequest.newBuilder(request) .setJobId(options.getJobId()) .setProjectId(options.getProject()) .setWorkerId(options.getWorkerId()) .build(), receiver); } GrpcWindmillServer(StreamingDataflowWorkerOptions options); private  GrpcWindmillServer(String name, boolean enableStreamingEngine); @Override synchronized void setWindmillServiceEndpoints(Set<HostAndPort> endpoints); @Override synchronized boolean isReady(); @Override void appendSummaryHtml(PrintWriter writer); @Override GetWorkResponse getWork(GetWorkRequest request); @Override GetDataResponse getData(GetDataRequest request); @Override CommitWorkResponse commitWork(CommitWorkRequest request); @Override GetWorkStream getWorkStream(GetWorkRequest request, WorkItemReceiver receiver); @Override GetDataStream getDataStream(); @Override CommitWorkStream commitWorkStream(); @Override GetConfigResponse getConfig(GetConfigRequest request); @Override ReportStatsResponse reportStats(ReportStatsRequest request); @Override long getAndResetThrottleTime();  }
@Test public void testStreamingGetWork() throws Exception { serviceRegistry.addService( new CloudWindmillServiceV1Alpha1ImplBase() { @Override public StreamObserver<StreamingGetWorkRequest> getWorkStream( StreamObserver<StreamingGetWorkResponseChunk> responseObserver) { return new StreamObserver<StreamingGetWorkRequest>() { boolean sawHeader = false; ResponseErrorInjector injector = new ResponseErrorInjector(responseObserver); @Override public void onNext(StreamingGetWorkRequest request) { maybeInjectError(responseObserver); try { long maxItems; if (!sawHeader) { errorCollector.checkThat( request.getRequest(), Matchers.equalTo( GetWorkRequest.newBuilder() .setClientId(10) .setJobId("job") .setProjectId("project") .setWorkerId("worker") .setMaxItems(3) .setMaxBytes(10000) .build())); sawHeader = true; maxItems = request.getRequest().getMaxItems(); } else { maxItems = request.getRequestExtension().getMaxItems(); } for (int item = 0; item < maxItems; item++) { long id = ThreadLocalRandom.current().nextLong(); ByteString serializedResponse = WorkItem.newBuilder() .setKey(ByteString.copyFromUtf8("somewhat_long_key")) .setWorkToken(id) .setShardingKey(id) .build() .toByteString(); for (int i = 0; i < serializedResponse.size(); i += 10) { int end = Math.min(serializedResponse.size(), i + 10); StreamingGetWorkResponseChunk.Builder builder = StreamingGetWorkResponseChunk.newBuilder() .setStreamId(id) .setSerializedWorkItem(serializedResponse.substring(i, end)) .setRemainingBytesForWorkItem(serializedResponse.size() - end); if (i == 0) { builder.setComputationMetadata( ComputationWorkItemMetadata.newBuilder() .setComputationId("comp") .setDependentRealtimeInputWatermark(17000) .setInputDataWatermark(18000)); } try { responseObserver.onNext(builder.build()); } catch (IllegalStateException e) { return; } } } } catch (Exception e) { errorCollector.addError(e); } } @Override public void onError(Throwable throwable) {} @Override public void onCompleted() { injector.cancel(); responseObserver.onCompleted(); } }; } }); CountDownLatch latch = new CountDownLatch(100); GetWorkStream stream = client.getWorkStream( GetWorkRequest.newBuilder().setClientId(10).setMaxItems(3).setMaxBytes(10000).build(), (String computation, @Nullable Instant inputDataWatermark, Instant synchronizedProcessingTime, Windmill.WorkItem workItem) -> { latch.countDown(); assertEquals(inputDataWatermark, new Instant(18)); assertEquals(synchronizedProcessingTime, new Instant(17)); assertEquals(workItem.getKey(), ByteString.copyFromUtf8("somewhat_long_key")); }); assertTrue(latch.await(30, TimeUnit.SECONDS)); stream.close(); assertTrue(stream.awaitTermination(30, TimeUnit.SECONDS)); }