@Override
  public <T extends State> T state(StateNamespace namespace, StateTag<T> address) {
    return workItemState.get(namespace, address, StateContexts.nullContext());
  }
@Test
  public void testMapViaMultimapEntriesAndKeysMergeLocalRemove() {
    final String tag = "map";
    StateTag<MapState<byte[], Integer>> addr =
        StateTags.map(tag, ByteArrayCoder.of(), VarIntCoder.of());
    MapState<byte[], Integer> mapState = underTestMapViaMultimap.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<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);

    ReadableState<Iterable<Map.Entry<byte[], Integer>>> entriesResult =
        mapState.entries().readLater();
    ReadableState<Iterable<byte[]>> keysResult = mapState.keys().readLater();
    waitAndSet(entriesFuture, Arrays.asList(multimapEntry(key1, 1), multimapEntry(key2, 2)), 30);
    waitAndSet(keysFuture, Arrays.asList(multimapEntry(key1), multimapEntry(key2)), 30);

    mapState.remove(dup(key1));
    mapState.put(key2, 8);
    mapState.put(dup(key3), 9);

    Iterable<Map.Entry<byte[], Integer>> entries = entriesResult.read();
    assertEquals(2, Iterables.size(entries));
    assertThat(
        entries,
        Matchers.containsInAnyOrder(multimapEntryMatcher(key2, 8), multimapEntryMatcher(key3, 9)));

    Iterable<byte[]> keys = keysResult.read();
    assertThat(keys, Matchers.containsInAnyOrder(key2, key3));
  }