@Override
  public <T extends State> T state(StateNamespace namespace, StateTag<T> address) {
    return workItemState.get(namespace, address, StateContexts.nullContext());
  }
@Test
  public void testMultimapBasicEntriesAndKeys() {
    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);

    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 =
        multimapState.entries().readLater();
    ReadableState<Iterable<byte[]>> keysResult = multimapState.keys().readLater();
    waitAndSet(
        entriesFuture,
        Arrays.asList(multimapEntry(key1, 1, 2, 3), multimapEntry(key2, 2, 3, 4)),
        30);
    waitAndSet(keysFuture, Arrays.asList(multimapEntry(key1), multimapEntry(key2)), 30);

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

    Iterable<byte[]> keys = keysResult.read();
    assertEquals(2, Iterables.size(keys));
    assertThat(keys, Matchers.containsInAnyOrder(key1, key2));
  }