@Override
  public K getKey() {
    return key;
  }
@Test
  public void testMapStateIterator() {
    final String stateId = "foo";
    final String countStateId = "count";

    DoFn<KV<String, KV<String, Integer>>, KV<String, Integer>> fn =
        new DoFn<KV<String, KV<String, Integer>>, KV<String, Integer>>() {

          @StateId(stateId)
          private final StateSpec<MapState<String, Integer>> mapState =
              StateSpecs.map(StringUtf8Coder.of(), VarIntCoder.of());

          @StateId(countStateId)
          private final StateSpec<CombiningState<Integer, int[], Integer>> countState =
              StateSpecs.combiningFromInputInternal(VarIntCoder.of(), Sum.ofIntegers());

          @ProcessElement
          public void processElement(
              ProcessContext c,
              @StateId(stateId) MapState<String, Integer> mapState,
              @StateId(countStateId) CombiningState<Integer, int[], Integer> count) {
            SamzaMapState<String, Integer> state = (SamzaMapState<String, Integer>) mapState;
            KV<String, Integer> value = c.element().getValue();
            state.put(value.getKey(), value.getValue());
            count.add(1);
            if (count.read() >= 4) {
              final List<KV<String, Integer>> content = new ArrayList<>();
              final Iterator<Map.Entry<String, Integer>> iterator = state.readIterator().read();
              while (iterator.hasNext()) {
                Map.Entry<String, Integer> entry = iterator.next();
                content.add(KV.of(entry.getKey(), entry.getValue()));
                c.output(KV.of(entry.getKey(), entry.getValue()));
              }

              assertEquals(
                  content, ImmutableList.of(KV.of("a", 97), KV.of("b", 42), KV.of("c", 12)));
            }
          }
        };

    PCollection<KV<String, Integer>> output =
        pipeline
            .apply(
                Create.of(
                    KV.of("hello", KV.of("a", 97)),
                    KV.of("hello", KV.of("b", 42)),
                    KV.of("hello", KV.of("b", 42)),
                    KV.of("hello", KV.of("c", 12))))
            .apply(ParDo.of(fn));

    PAssert.that(output).containsInAnyOrder(KV.of("a", 97), KV.of("b", 42), KV.of("c", 12));

    pipeline.run();
  }