WindmillStateReader { public Future<Instant> watermarkFuture(ByteString encodedTag, String stateFamily) { return stateFuture(new StateTag(StateTag.Kind.WATERMARK, encodedTag, stateFamily), null); } WindmillStateReader(
      MetricTrackingWindmillServerStub server,
      String computation,
      ByteString key,
      long shardingKey,
      long workToken); Future<Instant> watermarkFuture(ByteString encodedTag, String stateFamily); Future<T> valueFuture(ByteString encodedTag, String stateFamily, Coder<T> coder); Future<Iterable<T>> bagFuture(
      ByteString encodedTag, String stateFamily, Coder<T> elemCoder); void startBatchAndBlock(); long getBytesRead(); static final long MAX_BAG_BYTES; static final long MAX_KEY_BYTES; }
@Test public void testReadWatermark() throws Exception { Future<Instant> future = underTest.watermarkFuture(STATE_KEY_1, STATE_FAMILY); Mockito.verifyNoMoreInteractions(mockWindmill); Windmill.KeyedGetDataRequest.Builder expectedRequest = Windmill.KeyedGetDataRequest.newBuilder() .setKey(DATA_KEY) .setShardingKey(SHARDING_KEY) .setWorkToken(WORK_TOKEN) .setMaxBytes(WindmillStateReader.MAX_KEY_BYTES) .addWatermarkHoldsToFetch( Windmill.WatermarkHold.newBuilder() .setTag(STATE_KEY_1) .setStateFamily(STATE_FAMILY)); Windmill.KeyedGetDataResponse.Builder response = Windmill.KeyedGetDataResponse.newBuilder() .setKey(DATA_KEY) .addWatermarkHolds( Windmill.WatermarkHold.newBuilder() .setTag(STATE_KEY_1) .setStateFamily(STATE_FAMILY) .addTimestamps(5000000) .addTimestamps(6000000)); Mockito.when(mockWindmill.getStateData(COMPUTATION, expectedRequest.build())) .thenReturn(response.build()); Instant result = future.get(); Mockito.verify(mockWindmill).getStateData(COMPUTATION, expectedRequest.build()); assertThat(result, Matchers.equalTo(new Instant(5000))); assertNoReader(future); }
@Test public void testCachingWithinBatch() throws Exception { underTest.watermarkFuture(STATE_KEY_1, STATE_FAMILY); underTest.watermarkFuture(STATE_KEY_1, STATE_FAMILY); assertEquals(1, underTest.pendingLookups.size()); }