@Override
  public <T extends State> T state(StateNamespace namespace, StateTag<T> address) {
    return workItemState.get(namespace, address, StateContexts.nullContext());
  }
@Test
  public void testMultimapEntriesCombineCacheAndWindmill() {
    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[] key1 = "key1".getBytes(StandardCharsets.UTF_8);
    final byte[] key2 = "key2".getBytes(StandardCharsets.UTF_8);
    final byte[] key3 = "key3".getBytes(StandardCharsets.UTF_8);

    SettableFuture<Iterable<Integer>> entryFuture = SettableFuture.create();
    when(mockReader.multimapFetchSingleEntryFuture(
            encodeWithCoder(key1, ByteArrayCoder.of()),
            key(NAMESPACE, tag),
            STATE_FAMILY,
            VarIntCoder.of()))
        .thenReturn(entryFuture);

    // to set up the entry key1 as cache complete and add some local changes
    waitAndSet(entryFuture, weightedList(1, 2, 3), 30);
    multimapState.get(key1).read();
    multimapState.put(dup(key1), 2);
    multimapState.put(dup(key3), 20);

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

    // windmill contains extra entry key2, and this time the entries returned should not be cached.
    waitAndSet(
        entriesFuture,
        Arrays.asList(multimapEntry(key1, 1, 2, 3), multimapEntry(key2, 4, 5, 6)),
        30);
    waitAndSet(keysFuture, Arrays.asList(multimapEntry(key1), multimapEntry(key2)), 30);

    // key1 exist in both cache and windmill; key2 exists only in windmill; key3 exists only in
    // cache. They should all be merged.
    Iterable<Map.Entry<byte[], Integer>> entries = multimapState.entries().read();

    assertEquals(8, Iterables.size(entries));
    assertThat(
        entries,
        Matchers.containsInAnyOrder(
            multimapEntryMatcher(key1, 1),
            multimapEntryMatcher(key1, 2),
            multimapEntryMatcher(key1, 2),
            multimapEntryMatcher(key1, 3),
            multimapEntryMatcher(key2, 4),
            multimapEntryMatcher(key2, 5),
            multimapEntryMatcher(key2, 6),
            multimapEntryMatcher(key3, 20)));

    assertThat(multimapState.keys().read(), Matchers.containsInAnyOrder(key1, key2, key3));
  }