public static <K, V> Reshuffle<K, V> of() {
    return new Reshuffle<>();
  }
@Test
  @Category({ValidatesRunner.class})
  public void testAssignShardFn() {
    List<KV<String, Integer>> inputKvs = Lists.newArrayList();
    for (int i = 0; i < 10; i++) {
      inputKvs.addAll(ARBITRARY_KVS);
    }

    PCollection<KV<String, Integer>> input =
        pipeline.apply(
            Create.of(inputKvs).withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())));

    PCollection<Integer> output =
        input
            .apply(ParDo.of(new AssignShardFn<>(2)))
            .apply(GroupByKey.create())
            .apply(MapElements.into(integers()).via(KV::getKey));

    PAssert.that(output).containsInAnyOrder(ImmutableList.of(0, 1));

    pipeline.run();
  }