@Override
  public <T extends State> T state(StateNamespace namespace, StateTag<T> address) {
    return workItemState.get(namespace, address, StateContexts.nullContext());
  }
@Test
  public void testMultimapCacheComplete() {
    final String tag = "multimap";
    StateTag<MultimapState<byte[], Integer>> addr =
        StateTags.multimap(tag, ByteArrayCoder.of(), VarIntCoder.of());
    MultimapState<byte[], Integer> multimapState = underTest.state(NAMESPACE, addr);

    final byte[] key = "key".getBytes(StandardCharsets.UTF_8);

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

    // to set up the multimap as cache complete
    waitAndSet(entriesFuture, weightedList(multimapEntry(key, 1, 2, 3)), 30);
    multimapState.entries().read();

    multimapState.put(key, 2);

    when(mockReader.multimapFetchAllFuture(
            anyBoolean(), eq(key(NAMESPACE, tag)), eq(STATE_FAMILY), eq(VarIntCoder.of())))
        .thenThrow(
            new RuntimeException(
                "The multimap is cache complete and should not perform any windmill read."));
    when(mockReader.multimapFetchSingleEntryFuture(
            any(), eq(key(NAMESPACE, tag)), eq(STATE_FAMILY), eq(VarIntCoder.of())))
        .thenThrow(
            new RuntimeException(
                "The multimap is cache complete and should not perform any windmill read."));

    Iterable<Map.Entry<byte[], Integer>> entries = multimapState.entries().read();
    assertEquals(4, Iterables.size(entries));
    assertThat(
        entries,
        Matchers.containsInAnyOrder(
            multimapEntryMatcher(key, 1),
            multimapEntryMatcher(key, 2),
            multimapEntryMatcher(key, 3),
            multimapEntryMatcher(key, 2)));

    Iterable<byte[]> keys = multimapState.keys().read();
    assertThat(keys, Matchers.containsInAnyOrder(key));

    Iterable<Integer> values = multimapState.get(dup(key)).read();
    assertThat(values, Matchers.containsInAnyOrder(1, 2, 2, 3));
  }