@Override
  public <T extends State> T state(StateNamespace namespace, StateTag<T> address) {
    return workItemState.get(namespace, address, StateContexts.nullContext());
  }
@Test
  public void testMapAddBeforeRead() throws Exception {
    StateTag<MapState<String, Integer>> addr =
        StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of());
    MapState<String, Integer> mapState = underTest.state(NAMESPACE, addr);

    final String tag1 = "tag1";
    final String tag2 = "tag2";
    final String tag3 = "tag3";
    SettableFuture<Iterable<Map.Entry<ByteString, Integer>>> prefixFuture = SettableFuture.create();
    when(mockReader.valuePrefixFuture(
            protoKeyFromUserKey(null, StringUtf8Coder.of()), STATE_FAMILY, VarIntCoder.of()))
        .thenReturn(prefixFuture);

    ReadableState<Iterable<Map.Entry<String, Integer>>> result = mapState.entries();
    result = result.readLater();

    mapState.put(tag1, 1);
    waitAndSet(
        prefixFuture,
        ImmutableList.of(
            new AbstractMap.SimpleEntry<>(protoKeyFromUserKey(tag2, StringUtf8Coder.of()), 2)),
        200);
    Iterable<Map.Entry<String, Integer>> readData = result.read();
    assertThat(
        readData,
        Matchers.containsInAnyOrder(
            new AbstractMap.SimpleEntry<>(tag1, 1), new AbstractMap.SimpleEntry<>(tag2, 2)));

    mapState.put(tag3, 3);
    assertThat(
        result.read(),
        Matchers.containsInAnyOrder(
            new AbstractMap.SimpleEntry<>(tag1, 1),
            new AbstractMap.SimpleEntry<>(tag2, 2),
            new AbstractMap.SimpleEntry<>(tag3, 3)));
  }