public ForComputation forComputation(String computation) {
    return new ForComputation(computation);
  }
@Test
  public void testMultipleShardsOfKey() throws Exception {
    TestStateTag tag = new TestStateTag("tag1");

    WindmillStateCache.ForKeyAndFamily key1CacheShard1 =
        cache
            .forComputation(COMPUTATION)
            .forKey(computationKey(COMPUTATION, "key1", 1), 0L, 0L)
            .forFamily(STATE_FAMILY);
    WindmillStateCache.ForKeyAndFamily key1CacheShard2 =
        cache
            .forComputation(COMPUTATION)
            .forKey(computationKey(COMPUTATION, "key1", 2), 0L, 0L)
            .forFamily(STATE_FAMILY);
    WindmillStateCache.ForKeyAndFamily key2CacheShard1 =
        cache
            .forComputation(COMPUTATION)
            .forKey(computationKey(COMPUTATION, "key2", 1), 0L, 0L)
            .forFamily(STATE_FAMILY);

    TestState state1 = new TestState("g1");
    key1CacheShard1.put(StateNamespaces.global(), tag, state1, 2);
    key1CacheShard1.persist();
    assertEquals(Optional.of(state1), key1CacheShard1.get(StateNamespaces.global(), tag));
    key1CacheShard1 =
        cache
            .forComputation(COMPUTATION)
            .forKey(computationKey(COMPUTATION, "key1", 1), 0L, 1L)
            .forFamily(STATE_FAMILY);
    assertEquals(Optional.of(state1), key1CacheShard1.get(StateNamespaces.global(), tag));
    assertEquals(Optional.empty(), key1CacheShard2.get(StateNamespaces.global(), tag));
    assertEquals(Optional.empty(), key2CacheShard1.get(StateNamespaces.global(), tag));

    TestState state2 = new TestState("g2");
    key1CacheShard2.put(StateNamespaces.global(), tag, state2, 2);
    assertEquals(Optional.of(state2), key1CacheShard2.get(StateNamespaces.global(), tag));
    key1CacheShard2.persist();
    key1CacheShard2 =
        cache
            .forComputation(COMPUTATION)
            .forKey(computationKey(COMPUTATION, "key1", 2), 0L, 20L)
            .forFamily(STATE_FAMILY);
    assertEquals(Optional.of(state2), key1CacheShard2.get(StateNamespaces.global(), tag));
    assertEquals(Optional.of(state1), key1CacheShard1.get(StateNamespaces.global(), tag));
    assertEquals(Optional.empty(), key2CacheShard1.get(StateNamespaces.global(), tag));
  }