package org.apache.beam.runners.spark.structuredstreaming;

import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.Create;
import org.junit.BeforeClass;
import org.junit.Test;

/**
 * Test class for beam to spark source translation.
 */
public class SourceTest {
  private static Pipeline pipeline;

  @BeforeClass
  public static void beforeClass(){
    PipelineOptions options = PipelineOptionsFactory.create().as(SparkPipelineOptions.class);
    options.setRunner(SparkRunner.class);
    pipeline = Pipeline.create(options);
  }

  @Test
  public void testBoundedSource(){
    pipeline.apply(Create.of(1, 2, 3, 4, 5, 6, 7, 8, 9, 10));
    pipeline.run();
  }

}

@@ -53,16 +53,505 @@
 aming;%0A%0A
+import java.io.IOException;%0Aimport java.io.Serializable;%0Aimport java.util.ArrayList;%0Aimport java.util.HashMap;%0Aimport java.util.List;%0Aimport java.util.Map;%0Aimport org.apache.beam.runners.core.construction.PipelineOptionsSerializationUtils;%0Aimport org.apache.beam.runners.core.serialization.Base64Serializer;%0Aimport org.apache.beam.runners.spark.structuredstreaming.translation.batch.DatasetSourceBatch;%0Aimport org.apache.beam.runners.spark.structuredstreaming.utils.SerializationDebugger;%0A
 import o
@@ -579,16 +579,61 @@
 peline;%0A
+import org.apache.beam.sdk.io.BoundedSource;%0A
 import o
@@ -796,48 +796,299 @@
 org.
-junit.BeforeClass;%0Aimport org.junit.Test
+apache.spark.sql.sources.v2.DataSourceOptions;%0Aimport org.apache.spark.sql.sources.v2.reader.DataSourceReader;%0Aimport org.junit.BeforeClass;%0Aimport org.junit.Rule;%0Aimport org.junit.Test;%0Aimport org.junit.rules.TemporaryFolder;%0Aimport org.junit.runner.RunWith;%0Aimport org.junit.runners.JUnit4
 ;%0A%0A/
@@ -1146,16 +1146,39 @@
 on.%0A */%0A
+@RunWith(JUnit4.class)%0A
 public c
@@ -1193,16 +1193,40 @@
 rceTest 
+implements Serializable 
 %7B%0A  priv
@@ -1254,16 +1254,66 @@
 ipeline;
+%0A  @Rule%0A  public TemporaryFolder temporaryFolder;
 %0A%0A  @Bef
@@ -1539,16 +1539,1217 @@
 );%0A  %7D%0A%0A
+  @Test%0A  public void testSerialization() throws IOException%7B%0A    Map%3CString, String%3E datasetSourceOptions = new HashMap%3C%3E();%0A    BoundedSource%3CInteger%3E source = new BoundedSource%3CInteger%3E() %7B%0A%0A      @Override public List%3C? extends BoundedSource%3CInteger%3E%3E split(long desiredBundleSizeBytes,%0A          PipelineOptions options) throws Exception %7B%0A        return new ArrayList%3C%3E();%0A      %7D%0A%0A      @Override public long getEstimatedSizeBytes(PipelineOptions options) throws Exception %7B%0A        return 0;%0A      %7D%0A%0A      @Override public BoundedReader%3CInteger%3E createReader(PipelineOptions options)%0A          throws IOException %7B%0A        return null;%0A      %7D%0A    %7D;%0A    String serializedSource = Base64Serializer.serializeUnchecked(source);%0A    datasetSourceOptions.put(%22source%22, serializedSource);%0A    datasetSourceOptions.put(%22defaultParallelism%22, %224%22);%0A    datasetSourceOptions.put(%22pipelineOptions%22,%0A        PipelineOptionsSerializationUtils.serializeToJson(pipeline.getOptions()));%0A    DataSourceReader objectToTest = new DatasetSourceBatch()%0A        .createReader(new DataSourceOptions(datasetSourceOptions));%0A    SerializationDebugger.testSerialization(objectToTest, temporaryFolder.getRoot());%0A  %7D%0A%0A
   @Test%0A

Add serialization test