changed read test assertness and added another testDoFn

sdks/java/io/sparkreceiver/src/test/java/org/apache/beam/sdk/io/sparkreceiver/SparkReceiverIOTest.java

@@ -33,6 +33,8 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
+import java.util.List;
+
 /** Test class for {@link SparkReceiverIO}. */
 @RunWith(JUnit4.class)
 public class SparkReceiverIOTest {
@@ -113,10 +115,15 @@ public void testReadFromCustomReceiverWithOffset() {
  .withGetOffsetFn(Long::valueOf)
  .withSparkReceiverBuilder(receiverBuilder);
 
- TestOutputDoFn testDoFn = new TestOutputDoFn();
+ WithOffsetTestOutputDoFn testDoFn = new WithOffsetTestOutputDoFn();
+ List<String> storedRecords = CustomReceiverWithOffset.getStoredRecords();
+ List<String> outputRecords = WithOffsetTestOutputDoFn.getRecords();
+
  p.apply(reader).setCoder(StringUtf8Coder.of()).apply(ParDo.of(testDoFn));
  p.run().waitUntilFinish(Duration.standardSeconds(10));
- assertTrue(CustomReceiverWithOffset.getStoredRecords().containsAll(testDoFn.getRecords()));
+
+ List<String> expectedRecords = storedRecords.subList(0, storedRecords.size() - 2);
+ assertTrue(outputRecords.containsAll(expectedRecords));
  }
 
  @Test
@@ -135,9 +142,15 @@ public void testReadFromCustomReceiverWithoutOffset() {
  .withSparkConsumer(new CustomSparkConsumer<>())
  .withSparkReceiverBuilder(receiverBuilder);
 
- TestOutputDoFn testDoFn = new TestOutputDoFn();
+
+ WithoutOffsetTestOutputDoFn testDoFn = new WithoutOffsetTestOutputDoFn();
+ List<String> storedRecords = CustomReceiverWithoutOffset.getStoredRecords();
+ List<String> outputRecords = WithoutOffsetTestOutputDoFn.getRecords();
+
  p.apply(reader).setCoder(StringUtf8Coder.of()).apply(ParDo.of(testDoFn));
  p.run().waitUntilFinish(Duration.standardSeconds(10));
- assertTrue(CustomReceiverWithoutOffset.getStoredRecords().containsAll(testDoFn.getRecords()));
+
+ List<String> expectedRecords = storedRecords.subList(0, storedRecords.size() - 2);
+ assertTrue(outputRecords.containsAll(expectedRecords));
  }
 }

sdks/java/io/sparkreceiver/src/test/java/org/apache/beam/sdk/io/sparkreceiver/WithOffsetTestOutputDoFn.java

@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.sparkreceiver;
+
+import org.apache.beam.sdk.transforms.DoFn;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class WithOffsetTestOutputDoFn extends DoFn<String, String> {
+ private final static List<String> records = new ArrayList<>();
+
+ @ProcessElement
+ public void processElement(@Element String input, DoFn.OutputReceiver<String> output) {
+ records.add(input);
+ output.output(input);
+ }
+
+ public static List<String> getRecords() {
+ return records;
+ }
+}

sdks/java/io/sparkreceiver/src/test/java/org/apache/beam/sdk/io/sparkreceiver/TestOutputDoFn.java → sdks/java/io/sparkreceiver/src/test/java/org/apache/beam/sdk/io/sparkreceiver/WithoutOffsetTestOutputDoFn.java

@@ -21,16 +21,16 @@
 import java.util.List;
 import org.apache.beam.sdk.transforms.DoFn;
 
-public class TestOutputDoFn extends DoFn<String, String> {
- private static final List<String> records = new ArrayList<>();
+public class WithoutOffsetTestOutputDoFn extends DoFn<String, String> {
+ private final static List<String> records = new ArrayList<>();
 
  @ProcessElement
  public void processElement(@Element String input, OutputReceiver<String> output) {
  records.add(input);
  output.output(input);
  }
 
- public List<String> getRecords() {
+ public static List<String> getRecords() {
  return records;
  }
 }