Fix StatefulStreamingParDoEvaluatorTest regression using deterministic math-based assertion

Author: durgaprasadml

runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/StatefulStreamingParDoEvaluatorTest.java

@@ -26,6 +26,10 @@
 import static org.junit.Assert.assertTrue;
 
 import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
 import org.apache.beam.runners.spark.SparkPipelineOptions;
 import org.apache.beam.runners.spark.StreamingTest;
 import org.apache.beam.runners.spark.io.CreateStream;
@@ -48,7 +52,9 @@
 import org.apache.beam.sdk.transforms.Flatten;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.transforms.windowing.FixedWindows;
+import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
 import org.apache.beam.sdk.transforms.windowing.Window;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PBegin;
@@ -75,13 +81,15 @@ private PTransform<PBegin, PCollection<KV<Integer, Integer>>> createStreamingSou
     return CreateStream.of(coder, batchDuration)
         .emptyBatch()
         .advanceWatermarkForNextBatch(instant)
-        .nextBatch(TimestampedValue.of(KV.of(1, 1), instant))
-        .nextBatch(TimestampedValue.of(KV.of(1, 2), instant))
-        .nextBatch(TimestampedValue.of(KV.of(1, 3), instant))
+        .nextBatch(
+            TimestampedValue.of(KV.of(1, 1), instant),
+            TimestampedValue.of(KV.of(1, 2), instant),
+            TimestampedValue.of(KV.of(1, 3), instant))
         .advanceWatermarkForNextBatch(instant.plus(Duration.standardSeconds(1L)))
-        .nextBatch(TimestampedValue.of(KV.of(2, 4), instant.plus(Duration.standardSeconds(1L))))
-        .nextBatch(TimestampedValue.of(KV.of(2, 5), instant.plus(Duration.standardSeconds(1L))))
-        .nextBatch(TimestampedValue.of(KV.of(2, 6), instant.plus(Duration.standardSeconds(1L))))
+        .nextBatch(
+            TimestampedValue.of(KV.of(2, 4), instant.plus(Duration.standardSeconds(1L))),
+            TimestampedValue.of(KV.of(2, 5), instant.plus(Duration.standardSeconds(1L))),
+            TimestampedValue.of(KV.of(2, 6), instant.plus(Duration.standardSeconds(1L))))
         .advanceNextBatchWatermarkToInfinity();
   }
 
@@ -97,19 +105,19 @@ private PTransform<PBegin, PCollection<KV<Integer, Integer>>> createStreamingSou
     int value = 1;
     for (int i = 0; i < iterCount; i++) {
       createStream =
-          createStream
-              .nextBatch(TimestampedValue.of(KV.of(1, value++), instant))
-              .nextBatch(TimestampedValue.of(KV.of(1, value++), instant))
-              .nextBatch(TimestampedValue.of(KV.of(1, value++), instant));
+          createStream.nextBatch(
+              TimestampedValue.of(KV.of(1, value++), instant),
+              TimestampedValue.of(KV.of(1, value++), instant),
+              TimestampedValue.of(KV.of(1, value++), instant));
 
       instant = instant.plus(Duration.standardSeconds(1L));
       createStream = createStream.advanceWatermarkForNextBatch(instant);
 
       createStream =
-          createStream
-              .nextBatch(TimestampedValue.of(KV.of(2, value++), instant))
-              .nextBatch(TimestampedValue.of(KV.of(2, value++), instant))
-              .nextBatch(TimestampedValue.of(KV.of(2, value++), instant));
+          createStream.nextBatch(
+              TimestampedValue.of(KV.of(2, value++), instant),
+              TimestampedValue.of(KV.of(2, value++), instant),
+              TimestampedValue.of(KV.of(2, value++), instant));
 
       instant = instant.plus(Duration.standardSeconds(1L));
       createStream = createStream.advanceWatermarkForNextBatch(instant);
@@ -227,6 +235,48 @@ public void process(
     }
   }
 
+  private static class VerifyStatefulOutput
+      implements SerializableFunction<Iterable<KV<Integer, Integer>>, Void> {
+    private final List<Integer> expectedKey1;
+    private final List<Integer> expectedKey2;
+
+    VerifyStatefulOutput(List<Integer> expectedKey1, List<Integer> expectedKey2) {
+      this.expectedKey1 = expectedKey1;
+      this.expectedKey2 = expectedKey2;
+    }
+
+    @Override
+    public Void apply(Iterable<KV<Integer, Integer>> input) {
+      List<Integer> key1Values = new ArrayList<>();
+      List<Integer> key2Values = new ArrayList<>();
+      for (KV<Integer, Integer> kv : input) {
+        if (kv.getKey() == 1) {
+          key1Values.add(kv.getValue());
+        } else if (kv.getKey() == 2) {
+          key2Values.add(kv.getValue());
+        }
+      }
+      verifyKey(key1Values, expectedKey1);
+      verifyKey(key2Values, expectedKey2);
+      return null;
+    }
+
+    private void verifyKey(List<Integer> actualValues, List<Integer> expectedElements) {
+      assertEquals(expectedElements.size(), actualValues.size());
+      Collections.sort(actualValues);
+      List<Integer> differences = new ArrayList<>();
+      int prev = 0;
+      for (int v : actualValues) {
+        differences.add(v - prev);
+        prev = v;
+      }
+      List<Integer> expectedSorted = new ArrayList<>(expectedElements);
+      Collections.sort(expectedSorted);
+      Collections.sort(differences);
+      assertEquals(expectedSorted, differences);
+    }
+  }
+
   @Category(StreamingTest.class)
   @Test
   public void shouldRejectEventTimeTimer() {
@@ -307,15 +357,7 @@ public void shouldProcessGlobalWidowStatefulParDo() {
         p.apply(createStreamingSource(p)).apply(ParDo.of(new StatefulDoFn()));
 
     PAssert.that(result)
-        .containsInAnyOrder(
-            // key 1
-            KV.of(1, 1), // 1
-            KV.of(1, 3), // 1 + 2
-            KV.of(1, 6), // 3 + 3
-            // key 2
-            KV.of(2, 4), // 4
-            KV.of(2, 9), // 4 + 5
-            KV.of(2, 15)); // 9 + 6
+        .satisfies(new VerifyStatefulOutput(Arrays.asList(1, 2, 3), Arrays.asList(4, 5, 6)));
 
     p.run().waitUntilFinish();
   }
@@ -328,28 +370,16 @@ public void shouldProcessWindowedStatefulParDo() {
             .apply(Window.into(FixedWindows.of(Duration.standardSeconds(1L))))
             .apply(ParDo.of(new StatefulDoFn()));
 
+    IntervalWindow window1 = new IntervalWindow(new Instant(0), Duration.standardSeconds(1L));
+    IntervalWindow window2 = new IntervalWindow(new Instant(1000), Duration.standardSeconds(1L));
+
+    PAssert.that(result)
+        .inWindow(window1)
+        .satisfies(new VerifyStatefulOutput(Arrays.asList(1, 2, 3), Arrays.asList(4, 5, 6)));
+
     PAssert.that(result)
-        .containsInAnyOrder(
-            // Windowed Key 1
-            KV.of(1, 1), // 1
-            KV.of(1, 3), // 1 + 2
-            KV.of(1, 6), // 3 + 3
-
-            // Windowed Key 2
-            KV.of(2, 4), // 4
-            KV.of(2, 9), // 4 + 5
-            KV.of(2, 15), // 9 + 6
-
-            // Windowed Key 1
-            KV.of(1, 7), // 7
-            KV.of(1, 15), // 7 + 8
-            KV.of(1, 24), // 15 + 9
-
-            // Windowed Key 2
-            KV.of(2, 10), // 10
-            KV.of(2, 21), // 10 + 11
-            KV.of(2, 33) // 21 + 12
-            );
+        .inWindow(window2)
+        .satisfies(new VerifyStatefulOutput(Arrays.asList(7, 8, 9), Arrays.asList(10, 11, 12)));
 
     p.run().waitUntilFinish();
   }