Revert incorrect Spark stateful streaming test modifications

Author: durgaprasadml

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

@@ -26,10 +26,6 @@
 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;
@@ -52,9 +48,7 @@
 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;
@@ -235,48 +229,6 @@ 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() {
@@ -357,7 +309,15 @@ public void shouldProcessGlobalWidowStatefulParDo() {
         p.apply(createStreamingSource(p)).apply(ParDo.of(new StatefulDoFn()));
 
     PAssert.that(result)
-        .satisfies(new VerifyStatefulOutput(Arrays.asList(1, 2, 3), Arrays.asList(4, 5, 6)));
+        .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
 
     p.run().waitUntilFinish();
   }
@@ -370,26 +330,28 @@ 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));
-    IntervalWindow window3 = new IntervalWindow(new Instant(2000), Duration.standardSeconds(1L));
-    IntervalWindow window4 = new IntervalWindow(new Instant(3000), Duration.standardSeconds(1L));
-
-    PAssert.that(result)
-        .inWindow(window1)
-        .satisfies(new VerifyStatefulOutput(Arrays.asList(1, 2, 3), Collections.emptyList()));
-
-    PAssert.that(result)
-        .inWindow(window2)
-        .satisfies(new VerifyStatefulOutput(Collections.emptyList(), Arrays.asList(4, 5, 6)));
-
-    PAssert.that(result)
-        .inWindow(window3)
-        .satisfies(new VerifyStatefulOutput(Arrays.asList(7, 8, 9), Collections.emptyList()));
-
     PAssert.that(result)
-        .inWindow(window4)
-        .satisfies(new VerifyStatefulOutput(Collections.emptyList(), Arrays.asList(10, 11, 12)));
+        .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
+            );
 
     p.run().waitUntilFinish();
   }