WindmillStateReader { public <T> Future<Iterable<T>> bagFuture( ByteString encodedTag, String stateFamily, Coder<T> elemCoder) { StateTag stateTag = new StateTag(StateTag.Kind.BAG, encodedTag, stateFamily); return valuesToPagingIterableFuture( stateTag, elemCoder, this.<T, ValuesAndContPosition<T>>stateFuture(stateTag, elemCoder)); } 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 testReadBag() throws Exception { Future<Iterable<Integer>> future = underTest.bagFuture(STATE_KEY_1, STATE_FAMILY, INT_CODER); Mockito.verifyNoMoreInteractions(mockWindmill); Windmill.KeyedGetDataRequest.Builder expectedRequest = Windmill.KeyedGetDataRequest.newBuilder() .setKey(DATA_KEY) .setShardingKey(SHARDING_KEY) .setWorkToken(WORK_TOKEN) .setMaxBytes(WindmillStateReader.MAX_KEY_BYTES) .addBagsToFetch( Windmill.TagBag.newBuilder() .setTag(STATE_KEY_1) .setStateFamily(STATE_FAMILY) .setFetchMaxBytes(WindmillStateReader.MAX_BAG_BYTES)); Windmill.KeyedGetDataResponse.Builder response = Windmill.KeyedGetDataResponse.newBuilder() .setKey(DATA_KEY) .addBags( Windmill.TagBag.newBuilder() .setTag(STATE_KEY_1) .setStateFamily(STATE_FAMILY) .addValues(intData(5)) .addValues(intData(6))); Mockito.when(mockWindmill.getStateData(COMPUTATION, expectedRequest.build())) .thenReturn(response.build()); Iterable<Integer> results = future.get(); Mockito.verify(mockWindmill).getStateData(COMPUTATION, expectedRequest.build()); for (Integer unused : results) { } Mockito.verifyNoMoreInteractions(mockWindmill); assertThat(results, Matchers.contains(5, 6)); assertNoReader(future); }
@Test public void testReadBagWithContinuations() throws Exception { Future<Iterable<Integer>> future = underTest.bagFuture(STATE_KEY_1, STATE_FAMILY, INT_CODER); Mockito.verifyNoMoreInteractions(mockWindmill); Windmill.KeyedGetDataRequest.Builder expectedRequest1 = Windmill.KeyedGetDataRequest.newBuilder() .setKey(DATA_KEY) .setShardingKey(SHARDING_KEY) .setWorkToken(WORK_TOKEN) .setMaxBytes(WindmillStateReader.MAX_KEY_BYTES) .addBagsToFetch( Windmill.TagBag.newBuilder() .setTag(STATE_KEY_1) .setStateFamily(STATE_FAMILY) .setFetchMaxBytes(WindmillStateReader.MAX_BAG_BYTES)); Windmill.KeyedGetDataResponse.Builder response1 = Windmill.KeyedGetDataResponse.newBuilder() .setKey(DATA_KEY) .addBags( Windmill.TagBag.newBuilder() .setTag(STATE_KEY_1) .setStateFamily(STATE_FAMILY) .setContinuationPosition(CONT_POSITION) .addValues(intData(5)) .addValues(intData(6))); Windmill.KeyedGetDataRequest.Builder expectedRequest2 = Windmill.KeyedGetDataRequest.newBuilder() .setKey(DATA_KEY) .setShardingKey(SHARDING_KEY) .setWorkToken(WORK_TOKEN) .setMaxBytes(WindmillStateReader.MAX_KEY_BYTES) .addBagsToFetch( Windmill.TagBag.newBuilder() .setTag(STATE_KEY_1) .setStateFamily(STATE_FAMILY) .setFetchMaxBytes(WindmillStateReader.MAX_BAG_BYTES) .setRequestPosition(CONT_POSITION)); Windmill.KeyedGetDataResponse.Builder response2 = Windmill.KeyedGetDataResponse.newBuilder() .setKey(DATA_KEY) .addBags( Windmill.TagBag.newBuilder() .setTag(STATE_KEY_1) .setStateFamily(STATE_FAMILY) .setRequestPosition(CONT_POSITION) .addValues(intData(7)) .addValues(intData(8))); Mockito.when(mockWindmill.getStateData(COMPUTATION, expectedRequest1.build())) .thenReturn(response1.build()); Mockito.when(mockWindmill.getStateData(COMPUTATION, expectedRequest2.build())) .thenReturn(response2.build()); Iterable<Integer> results = future.get(); Mockito.verify(mockWindmill).getStateData(COMPUTATION, expectedRequest1.build()); for (Integer unused : results) { } Mockito.verify(mockWindmill).getStateData(COMPUTATION, expectedRequest2.build()); Mockito.verifyNoMoreInteractions(mockWindmill); assertThat(results, Matchers.contains(5, 6, 7, 8)); }