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

    SettableFuture<Iterable<Map.Entry<ByteString, Iterable<Integer>>>> keysFuture =
        SettableFuture.create();
    when(mockReader.multimapFetchAllFuture(
            true, key(NAMESPACE, tag), STATE_FAMILY, VarIntCoder.of()))
        .thenReturn(keysFuture);

    waitAndSet(
        keysFuture,
        () ->
            new Iterator<Map.Entry<ByteString, Iterable<Integer>>>() {
              final int targetEntries = 1_000_000; // return 1 million entries, which is 10 GBs
              final byte[] entryKey = new byte[10_000]; // each key is 10KB
              final Random rand = new Random();
              int returnedEntries = 0;

              @Override
              public boolean hasNext() {
                return returnedEntries < targetEntries;
              }

              @Override
              public Map.Entry<ByteString, Iterable<Integer>> next() {
                returnedEntries++;
                rand.nextBytes(entryKey);
                return multimapEntry(entryKey);
              }
            },
        200);
    Iterable<byte[]> keys = multimapState.keys().read();
    assertEquals(1_000_000, Iterables.size(keys));
  }