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

    final byte[] key = "key".getBytes(StandardCharsets.UTF_8);
    SettableFuture<Iterable<Integer>> future = SettableFuture.create();
    when(mockReader.multimapFetchSingleEntryFuture(
            encodeWithCoder(key, ByteArrayCoder.of()),
            key(NAMESPACE, tag),
            STATE_FAMILY,
            VarIntCoder.of()))
        .thenReturn(future);

    mapViaMultiMapState.put(key, 1);
    ReadableState<Integer> result = mapViaMultiMapState.get(dup(key)).readLater();
    waitAndSet(future, Collections.singletonList(2), 30);
    assertEquals(Integer.valueOf(1), result.read());

    mapViaMultiMapState.put(key, 3);
    assertEquals(Integer.valueOf(3), mapViaMultiMapState.get(key).read());
    mapViaMultiMapState.clear();
    assertNull(mapViaMultiMapState.get(key).read());
  }