@Override
  public <T extends State> T state(StateNamespace namespace, StateTag<T> address) {
    return workItemState.get(namespace, address, StateContexts.nullContext());
  }
@Test
  public void testMultimapLazyIterateHugeEntriesResultSingleEntry() {
    // A multimap with 1 key and 1 million values and a total of 10GBs data
    final String tag = "multimap";
    final Integer key = 100;
    StateTag<MultimapState<Integer, byte[]>> addr =
        StateTags.multimap(tag, VarIntCoder.of(), ByteArrayCoder.of());
    MultimapState<Integer, byte[]> multimapState = underTest.state(NAMESPACE, addr);

    SettableFuture<Iterable<Map.Entry<ByteString, Iterable<byte[]>>>> entriesFuture =
        SettableFuture.create();
    when(mockReader.multimapFetchAllFuture(
            false, key(NAMESPACE, tag), STATE_FAMILY, ByteArrayCoder.of()))
        .thenReturn(entriesFuture);
    SettableFuture<Iterable<byte[]>> getKeyFuture = SettableFuture.create();
    when(mockReader.multimapFetchSingleEntryFuture(
            encodeWithCoder(key, VarIntCoder.of()),
            key(NAMESPACE, tag),
            STATE_FAMILY,
            ByteArrayCoder.of()))
        .thenReturn(getKeyFuture);

    // a not weighted iterators that returns tons of data
    Iterable<byte[]> values =
        () ->
            new Iterator<byte[]>() {
              final int targetValues = 1_000_000; // return 1 million values, which is 10 GBs
              final byte[] value = new byte[10_000]; // each value is 10KB
              final Random rand = new Random();
              int returnedValues = 0;

              @Override
              public boolean hasNext() {
                return returnedValues < targetValues;
              }

              @Override
              public byte[] next() {
                returnedValues++;
                rand.nextBytes(value);
                return value;
              }
            };

    waitAndSet(
        entriesFuture,
        Collections.singletonList(
            new SimpleEntry<>(encodeWithCoder(key, VarIntCoder.of()), values)),
        200);
    waitAndSet(getKeyFuture, values, 200);

    Iterable<Map.Entry<Integer, byte[]>> entries = multimapState.entries().read();
    assertEquals(1_000_000, Iterables.size(entries));

    Iterable<byte[]> valueResult = multimapState.get(key).read();
    assertEquals(1_000_000, Iterables.size(valueResult));
  }