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

import java.io.IOException;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.apache.beam.runners.core.construction.PipelineOptionsSerializationUtils;
import org.apache.beam.runners.core.serialization.Base64Serializer;
import org.apache.beam.runners.spark.structuredstreaming.SparkPipelineOptions;
import org.apache.beam.runners.spark.structuredstreaming.SparkRunner;
import org.apache.beam.runners.spark.structuredstreaming.translation.batch.DatasetSourceBatch;
import org.apache.beam.runners.spark.structuredstreaming.utils.SerializationDebugger;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.BoundedSource;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.Create;
import org.apache.spark.sql.sources.v2.DataSourceOptions;
import org.apache.spark.sql.sources.v2.reader.DataSourceReader;
import org.junit.BeforeClass;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

/**
 * Test class for beam to spark source translation.
 */
@RunWith(JUnit4.class)
public class SourceTest implements Serializable {
  private static Pipeline pipeline;
  @Rule
  public TemporaryFolder temporaryFolder;

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

  @Test
  public void testSerialization() throws IOException{
    Map<String, String> datasetSourceOptions = new HashMap<>();
    BoundedSource<Integer> source = new BoundedSource<Integer>() {

      @Override public List<? extends BoundedSource<Integer>> split(long desiredBundleSizeBytes,
          PipelineOptions options) throws Exception {
        return new ArrayList<>();
      }

      @Override public long getEstimatedSizeBytes(PipelineOptions options) throws Exception {
        return 0;
      }

      @Override public BoundedReader<Integer> createReader(PipelineOptions options)
          throws IOException {
        return null;
      }
    };
    String serializedSource = Base64Serializer.serializeUnchecked(source);
    datasetSourceOptions.put("source", serializedSource);
    datasetSourceOptions.put("defaultParallelism", "4");
    datasetSourceOptions.put("pipelineOptions",
        PipelineOptionsSerializationUtils.serializeToJson(pipeline.getOptions()));
    DataSourceReader objectToTest = new DatasetSourceBatch()
        .createReader(new DataSourceOptions(datasetSourceOptions));
    SerializationDebugger.testSerialization(objectToTest, temporaryFolder.getRoot());
  }

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

}

@@ -532,103 +532,8 @@
 er;%0A
-import org.apache.beam.runners.spark.structuredstreaming.translation.batch.DatasetSourceBatch;%0A
 impo
@@ -1022,16 +1022,21 @@
 g.junit.
+Class
 Rule;%0Aim
@@ -1335,16 +1335,21 @@
 ine;%0A  @
+Class
 Rule%0A  p
@@ -1354,16 +1354,29 @@
  public 
+static final 
 Temporar
@@ -1383,17 +1383,40 @@
 yFolder 
-t
+TEMPORARY_FOLDER = new T
 emporary
@@ -1421,16 +1421,18 @@
 ryFolder
+()
 ;%0A%0A  @Be
@@ -1725,72 +1725,8 @@
 on%7B%0A
-    Map%3CString, String%3E datasetSourceOptions = new HashMap%3C%3E();%0A
     
@@ -2338,41 +2338,128 @@
     
-datasetSourceOptions.put(%22source%22
+Map%3CString, String%3E dataSourceOptions = new HashMap%3C%3E();%0A    dataSourceOptions.put(DatasetSourceBatch.BEAM_SOURCE_OPTION
 , se
@@ -2475,35 +2475,32 @@
 ource);%0A    data
-set
 SourceOptions.pu
@@ -2505,28 +2505,46 @@
 put(
-%22defaultParallelism%22
+DatasetSourceBatch.DEFAULT_PARALLELISM
 , %224
@@ -2551,27 +2551,24 @@
 %22);%0A    data
-set
 SourceOption
@@ -2577,25 +2577,43 @@
 put(
-%22pipelineOptions%22
+DatasetSourceBatch.PIPELINE_OPTIONS
 ,%0A  
@@ -2802,19 +2802,16 @@
 ons(data
-set
 SourceOp
@@ -2881,31 +2881,32 @@
 st, 
-temporaryFolder.getRoot
+TEMPORARY_FOLDER.newFile
 ());

Fix SourceTest