made adjustments to read assertions

sdks/java/io/sparkreceiver/src/main/java/org/apache/beam/sdk/io/sparkreceiver/ReadFromSparkReceiverWithOffsetDoFn.java

@@ -148,7 +148,6 @@ public void start(Receiver<V> sparkReceiver) {
  }
  ((HasOffset) sparkReceiver).setStartOffset(startOffset);
  sparkReceiver.supervisor().startReceiver();
- LOG.info("Spark receiver was started");
  try {
  TimeUnit.MILLISECONDS.sleep(START_POLL_TIMEOUT_MS);
  } catch (InterruptedException e) {

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

@@ -67,8 +67,10 @@ public Long getEndOffset() {
  private void receive() {
  Long currentOffset = startOffset;
  while (!isStopped()) {
- STORED_RECORDS.add(currentOffset.toString());
- store((currentOffset++).toString());
+ if (currentOffset <= 20) {
+ STORED_RECORDS.add(currentOffset.toString());
+ store((currentOffset++).toString());
+ }
  try {
  TimeUnit.MILLISECONDS.sleep(TIMEOUT_MS);
  } catch (InterruptedException e) {

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

@@ -53,8 +53,10 @@ public void onStop() {}
  private void receive() {
  Long currentOffset = 0L;
  while (!isStopped()) {
- STORED_RECORDS.add(currentOffset.toString());
- store((currentOffset++).toString());
+ if (currentOffset <= 20) {
+ STORED_RECORDS.add(currentOffset.toString());
+ store((currentOffset++).toString());
+ }
  try {
  TimeUnit.MILLISECONDS.sleep(TIMEOUT_MS);
  } catch (InterruptedException e) {

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

@@ -19,7 +19,6 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThrows;
-import static org.junit.Assert.assertTrue;
 
 import java.util.List;
 import org.apache.beam.runners.direct.DirectOptions;
@@ -114,15 +113,14 @@ public void testReadFromCustomReceiverWithOffset() {
  .withGetOffsetFn(Long::valueOf)
  .withSparkReceiverBuilder(receiverBuilder);
 
- WithOffsetTestOutputDoFn testDoFn = new WithOffsetTestOutputDoFn();
  List<String> storedRecords = CustomReceiverWithOffset.getStoredRecords();
- List<String> outputRecords = WithOffsetTestOutputDoFn.getRecords();
+ List<String> outputRecords = TestOutputDoFn.getRecords();
+ outputRecords.clear();
 
- p.apply(reader).setCoder(StringUtf8Coder.of()).apply(ParDo.of(testDoFn));
- p.run().waitUntilFinish(Duration.standardSeconds(10));
+ p.apply(reader).setCoder(StringUtf8Coder.of()).apply(ParDo.of(new TestOutputDoFn()));
+ p.run().waitUntilFinish(Duration.standardSeconds(15));
 
- List<String> expectedRecords = storedRecords.subList(0, storedRecords.size() - 2);
- assertTrue(outputRecords.containsAll(expectedRecords));
+ assertEquals(outputRecords, storedRecords);
  }
 
  @Test
@@ -141,14 +139,13 @@ public void testReadFromCustomReceiverWithoutOffset() {
  .withSparkConsumer(new CustomSparkConsumer<>())
  .withSparkReceiverBuilder(receiverBuilder);
 
- WithoutOffsetTestOutputDoFn testDoFn = new WithoutOffsetTestOutputDoFn();
  List<String> storedRecords = CustomReceiverWithoutOffset.getStoredRecords();
- List<String> outputRecords = WithoutOffsetTestOutputDoFn.getRecords();
+ List<String> outputRecords = TestOutputDoFn.getRecords();
+ outputRecords.clear();
 
- p.apply(reader).setCoder(StringUtf8Coder.of()).apply(ParDo.of(testDoFn));
- p.run().waitUntilFinish(Duration.standardSeconds(10));
+ p.apply(reader).setCoder(StringUtf8Coder.of()).apply(ParDo.of(new TestOutputDoFn()));
+ p.run().waitUntilFinish(Duration.standardSeconds(15));
 
- List<String> expectedRecords = storedRecords.subList(0, storedRecords.size() - 2);
- assertTrue(outputRecords.containsAll(expectedRecords));
+ assertEquals(storedRecords, outputRecords);
  }
 }

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

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