beam-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From amits...@apache.org
Subject [18/23] beam git commit: Use a PipelineRule for test pipelines.
Date Tue, 28 Feb 2017 22:35:25 GMT
Use a PipelineRule for test pipelines.


Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/4ca56806
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/4ca56806
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/4ca56806

Branch: refs/heads/master
Commit: 4ca5680681e14c4be8b0ac1f3e1244c9d439e6d8
Parents: bd41d9a
Author: Sela <ansela@paypal.com>
Authored: Mon Feb 27 19:24:45 2017 +0200
Committer: Sela <ansela@paypal.com>
Committed: Wed Mar 1 00:18:07 2017 +0200

----------------------------------------------------------------------
 runners/spark/pom.xml                           |   1 -
 .../runners/spark/SparkPipelineOptions.java     |   9 --
 .../runners/spark/TestSparkPipelineOptions.java |  18 ++++
 .../beam/runners/spark/TestSparkRunner.java     |  28 ++---
 .../beam/runners/spark/ForceStreamingTest.java  |   8 +-
 .../apache/beam/runners/spark/PipelineRule.java | 103 +++++++++++++++++++
 .../metrics/sink/NamedAggregatorsTest.java      |   9 +-
 .../beam/runners/spark/io/AvroPipelineTest.java |   6 +-
 .../beam/runners/spark/io/NumShardsTest.java    |   6 +-
 .../io/hadoop/HadoopFileFormatPipelineTest.java |   6 +-
 .../spark/translation/StorageLevelTest.java     |   8 +-
 .../translation/streaming/CreateStreamTest.java |  77 ++++----------
 .../utils/SparkTestPipelineOptions.java         |  42 --------
 .../SparkTestPipelineOptionsForStreaming.java   |  48 ---------
 14 files changed, 175 insertions(+), 194 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/beam/blob/4ca56806/runners/spark/pom.xml
----------------------------------------------------------------------
diff --git a/runners/spark/pom.xml b/runners/spark/pom.xml
index f926bf5..409fc27 100644
--- a/runners/spark/pom.xml
+++ b/runners/spark/pom.xml
@@ -80,7 +80,6 @@
                     org.apache.beam.sdk.testing.UsesCommittedMetrics
                   </excludedGroups>
                   <parallel>none</parallel>
-                  <threadCount>4</threadCount>
                   <forkCount>1</forkCount>
                   <reuseForks>false</reuseForks>
                   <failIfNoTests>true</failIfNoTests>

http://git-wip-us.apache.org/repos/asf/beam/blob/4ca56806/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineOptions.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineOptions.java
b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineOptions.java
index 52d8ce1..26b549b 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineOptions.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineOptions.java
@@ -101,13 +101,4 @@ public interface SparkPipelineOptions
   boolean getUsesProvidedSparkContext();
   void setUsesProvidedSparkContext(boolean value);
 
-  @Description("A special flag that forces streaming in tests.")
-  @Default.Boolean(false)
-  boolean isForceStreaming();
-  void setForceStreaming(boolean forceStreaming);
-
-  @Description("A forced timeout (millis), mostly for testing.")
-  @Default.Long(3000L)
-  Long getForcedTimeout();
-  void setForcedTimeout(Long forcedTimeout);
 }

http://git-wip-us.apache.org/repos/asf/beam/blob/4ca56806/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkPipelineOptions.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkPipelineOptions.java
b/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkPipelineOptions.java
new file mode 100644
index 0000000..2cb58d8
--- /dev/null
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkPipelineOptions.java
@@ -0,0 +1,18 @@
+package org.apache.beam.runners.spark;
+
+import org.apache.beam.sdk.options.Default;
+import org.apache.beam.sdk.options.Description;
+import org.apache.beam.sdk.testing.TestPipelineOptions;
+
+
+/**
+ * A {@link SparkPipelineOptions} for tests.
+ */
+public interface TestSparkPipelineOptions extends SparkPipelineOptions, TestPipelineOptions
{
+
+  @Description("A special flag that forces streaming in tests.")
+  @Default.Boolean(false)
+  boolean isForceStreaming();
+  void setForceStreaming(boolean forceStreaming);
+
+}

http://git-wip-us.apache.org/repos/asf/beam/blob/4ca56806/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkRunner.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkRunner.java
b/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkRunner.java
index 5d71ea5..035da00 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkRunner.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkRunner.java
@@ -18,6 +18,7 @@
 
 package org.apache.beam.runners.spark;
 
+import static com.google.common.base.Preconditions.checkNotNull;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.is;
 
@@ -33,7 +34,6 @@ import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.options.PipelineOptionsValidator;
 import org.apache.beam.sdk.runners.PipelineRunner;
 import org.apache.beam.sdk.testing.PAssert;
-import org.apache.beam.sdk.testing.TestPipelineOptions;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.util.ValueWithRecordId;
@@ -76,15 +76,15 @@ public final class TestSparkRunner extends PipelineRunner<SparkPipelineResult>
{
   private boolean isForceStreaming;
   private int expectedNumberOfAssertions = 0;
 
-  private TestSparkRunner(SparkPipelineOptions options) {
+  private TestSparkRunner(TestSparkPipelineOptions options) {
     this.delegate = SparkRunner.fromOptions(options);
     this.isForceStreaming = options.isForceStreaming();
   }
 
   public static TestSparkRunner fromOptions(PipelineOptions options) {
     // Default options suffice to set it up as a test runner
-    SparkPipelineOptions sparkOptions =
-        PipelineOptionsValidator.validate(SparkPipelineOptions.class, options);
+    TestSparkPipelineOptions sparkOptions =
+        PipelineOptionsValidator.validate(TestSparkPipelineOptions.class, options);
     return new TestSparkRunner(sparkOptions);
   }
 
@@ -115,22 +115,22 @@ public final class TestSparkRunner extends PipelineRunner<SparkPipelineResult>
{
 
   @Override
   public SparkPipelineResult run(Pipeline pipeline) {
-    SparkPipelineOptions sparkOptions = pipeline.getOptions().as(SparkPipelineOptions.class);
+    TestSparkPipelineOptions testSparkPipelineOptions =
+        pipeline.getOptions().as(TestSparkPipelineOptions.class);
     SparkPipelineResult result = null;
     // clear state of Aggregators, Metrics and Watermarks.
     AggregatorsAccumulator.clear();
     SparkMetricsContainer.clear();
     GlobalWatermarkHolder.clear();
 
-    TestPipelineOptions testPipelineOptions = pipeline.getOptions().as(TestPipelineOptions.class);
-    LOG.info("About to run test pipeline " + sparkOptions.getJobName());
+    LOG.info("About to run test pipeline " + testSparkPipelineOptions.getJobName());
 
     // if the pipeline was executed in streaming mode, validate aggregators.
     if (isForceStreaming) {
       try {
         result = delegate.run(pipeline);
-        long timeout = sparkOptions.getForcedTimeout();
-        result.waitUntilFinish(Duration.millis(timeout));
+        Long timeout = testSparkPipelineOptions.getTestTimeoutSeconds();
+        result.waitUntilFinish(Duration.standardSeconds(checkNotNull(timeout)));
         // validate assertion succeeded (at least once).
         int successAssertions = result.getAggregatorValue(PAssert.SUCCESS_COUNTER, Integer.class);
         assertThat(
@@ -149,12 +149,12 @@ public final class TestSparkRunner extends PipelineRunner<SparkPipelineResult>
{
         LOG.info(
             String.format(
                 "Successfully asserted pipeline %s with %d successful assertions.",
-                sparkOptions.getJobName(),
+                testSparkPipelineOptions.getJobName(),
                 successAssertions));
-        } finally {
+      } finally {
         try {
           // cleanup checkpoint dir.
-          FileUtils.deleteDirectory(new File(sparkOptions.getCheckpointDir()));
+          FileUtils.deleteDirectory(new File(testSparkPipelineOptions.getCheckpointDir()));
         } catch (IOException e) {
           throw new RuntimeException("Failed to clear checkpoint tmp dir.", e);
         }
@@ -164,8 +164,8 @@ public final class TestSparkRunner extends PipelineRunner<SparkPipelineResult>
{
       result = delegate.run(pipeline);
       result.waitUntilFinish();
       // assert via matchers.
-      assertThat(result, testPipelineOptions.getOnCreateMatcher());
-      assertThat(result, testPipelineOptions.getOnSuccessMatcher());
+      assertThat(result, testSparkPipelineOptions.getOnCreateMatcher());
+      assertThat(result, testSparkPipelineOptions.getOnSuccessMatcher());
     }
     return result;
   }

http://git-wip-us.apache.org/repos/asf/beam/blob/4ca56806/runners/spark/src/test/java/org/apache/beam/runners/spark/ForceStreamingTest.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/ForceStreamingTest.java
b/runners/spark/src/test/java/org/apache/beam/runners/spark/ForceStreamingTest.java
index c3026ce..9b39558 100644
--- a/runners/spark/src/test/java/org/apache/beam/runners/spark/ForceStreamingTest.java
+++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/ForceStreamingTest.java
@@ -21,7 +21,6 @@ package org.apache.beam.runners.spark;
 import static org.hamcrest.MatcherAssert.assertThat;
 
 import java.io.IOException;
-import org.apache.beam.runners.spark.translation.streaming.utils.SparkTestPipelineOptionsForStreaming;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.io.BoundedReadFromUnboundedSource;
 import org.apache.beam.sdk.io.CountingSource;
@@ -46,14 +45,11 @@ import org.junit.Test;
 public class ForceStreamingTest {
 
   @Rule
-  public SparkTestPipelineOptionsForStreaming commonOptions =
-      new SparkTestPipelineOptionsForStreaming();
+  public final PipelineRule pipelineRule = PipelineRule.streaming();
 
   @Test
   public void test() throws IOException {
-    SparkPipelineOptions options = commonOptions.getOptions();
-    options.setForceStreaming(true);
-    Pipeline pipeline = Pipeline.create(options);
+    Pipeline pipeline = pipelineRule.createPipeline();
 
     // apply the BoundedReadFromUnboundedSource.
     BoundedReadFromUnboundedSource<?> boundedRead =

http://git-wip-us.apache.org/repos/asf/beam/blob/4ca56806/runners/spark/src/test/java/org/apache/beam/runners/spark/PipelineRule.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/PipelineRule.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/PipelineRule.java
new file mode 100644
index 0000000..bb42510
--- /dev/null
+++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/PipelineRule.java
@@ -0,0 +1,103 @@
+package org.apache.beam.runners.spark;
+
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.joda.time.Duration;
+import org.junit.rules.ExternalResource;
+import org.junit.rules.RuleChain;
+import org.junit.rules.TemporaryFolder;
+import org.junit.rules.TestName;
+import org.junit.rules.TestRule;
+import org.junit.runner.Description;
+import org.junit.runners.model.Statement;
+
+/**
+ * A {@link org.junit.Rule} to provide a {@link Pipeline} instance for Spark runner tests.
+ */
+public class PipelineRule implements TestRule {
+
+  private final TestName testName = new TestName();
+
+  private final SparkPipelineRule delegate;
+  private final RuleChain chain;
+
+  private PipelineRule() {
+    this.delegate = new SparkPipelineRule(testName);
+    this.chain = RuleChain.outerRule(testName).around(this.delegate);
+  }
+
+  private PipelineRule(Duration forcedTimeout) {
+    this.delegate = new SparkStreamingPipelineRule(forcedTimeout, testName);
+    this.chain = RuleChain.outerRule(testName).around(this.delegate);
+  }
+
+  public static PipelineRule streaming() {
+    return new PipelineRule(Duration.standardSeconds(5));
+  }
+
+  public static PipelineRule batch() {
+    return new PipelineRule();
+  }
+
+  public Duration batchDuration() {
+    return Duration.millis(delegate.options.getBatchIntervalMillis());
+  }
+
+  public SparkPipelineOptions getOptions() {
+    return delegate.options;
+  }
+
+  public Pipeline createPipeline() {
+    return Pipeline.create(delegate.options);
+  }
+
+  @Override
+  public Statement apply(Statement statement, Description description) {
+    return chain.apply(statement, description);
+  }
+
+  private static class SparkStreamingPipelineRule extends SparkPipelineRule {
+
+    private final TemporaryFolder temporaryFolder = new TemporaryFolder();
+    private final Duration forcedTimeout;
+
+    SparkStreamingPipelineRule(Duration forcedTimeout, TestName testName) {
+      super(testName);
+      this.forcedTimeout = forcedTimeout;
+    }
+
+    @Override
+    protected void before() throws Throwable {
+      super.before();
+      temporaryFolder.create();
+      options.setForceStreaming(true);
+      options.setTestTimeoutSeconds(forcedTimeout.getStandardSeconds());
+      options.setCheckpointDir(
+          temporaryFolder.newFolder(options.getJobName()).toURI().toURL().toString());
+    }
+
+    @Override
+    protected void after() {
+      temporaryFolder.delete();
+    }
+  }
+
+  private static class SparkPipelineRule extends ExternalResource {
+
+    protected final TestSparkPipelineOptions options =
+        PipelineOptionsFactory.as(TestSparkPipelineOptions.class);
+
+    private final TestName testName;
+
+    private SparkPipelineRule(TestName testName) {
+      this.testName = testName;
+    }
+
+    @Override
+    protected void before() throws Throwable {
+      options.setRunner(TestSparkRunner.class);
+      options.setEnableSparkMetricSinks(false);
+      options.setJobName(testName.getMethodName());
+    }
+  }
+}

http://git-wip-us.apache.org/repos/asf/beam/blob/4ca56806/runners/spark/src/test/java/org/apache/beam/runners/spark/aggregators/metrics/sink/NamedAggregatorsTest.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/aggregators/metrics/sink/NamedAggregatorsTest.java
b/runners/spark/src/test/java/org/apache/beam/runners/spark/aggregators/metrics/sink/NamedAggregatorsTest.java
index 2f7202c..a192807 100644
--- a/runners/spark/src/test/java/org/apache/beam/runners/spark/aggregators/metrics/sink/NamedAggregatorsTest.java
+++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/aggregators/metrics/sink/NamedAggregatorsTest.java
@@ -26,12 +26,12 @@ import com.google.common.collect.ImmutableSet;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Set;
+import org.apache.beam.runners.spark.PipelineRule;
 import org.apache.beam.runners.spark.SparkPipelineOptions;
 import org.apache.beam.runners.spark.aggregators.ClearAggregatorsRule;
 import org.apache.beam.runners.spark.aggregators.SparkAggregators;
 import org.apache.beam.runners.spark.examples.WordCount;
 import org.apache.beam.runners.spark.translation.SparkContextFactory;
-import org.apache.beam.runners.spark.translation.streaming.utils.SparkTestPipelineOptions;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
@@ -56,12 +56,11 @@ public class NamedAggregatorsTest {
   public ClearAggregatorsRule clearAggregators = new ClearAggregatorsRule();
 
   @Rule
-  public final SparkTestPipelineOptions pipelineOptions = new SparkTestPipelineOptions();
+  public final PipelineRule pipelineRule = PipelineRule.batch();
 
   private Pipeline createSparkPipeline() {
-    SparkPipelineOptions options = pipelineOptions.getOptions();
-    options.setEnableSparkMetricSinks(true);
-    return Pipeline.create(options);
+    pipelineRule.getOptions().setEnableSparkMetricSinks(true);
+    return pipelineRule.createPipeline();
   }
 
   private void runPipeline() {

http://git-wip-us.apache.org/repos/asf/beam/blob/4ca56806/runners/spark/src/test/java/org/apache/beam/runners/spark/io/AvroPipelineTest.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/io/AvroPipelineTest.java
b/runners/spark/src/test/java/org/apache/beam/runners/spark/io/AvroPipelineTest.java
index c5bb583..2a73c28 100644
--- a/runners/spark/src/test/java/org/apache/beam/runners/spark/io/AvroPipelineTest.java
+++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/io/AvroPipelineTest.java
@@ -33,7 +33,7 @@ import org.apache.avro.generic.GenericData;
 import org.apache.avro.generic.GenericDatumReader;
 import org.apache.avro.generic.GenericDatumWriter;
 import org.apache.avro.generic.GenericRecord;
-import org.apache.beam.runners.spark.translation.streaming.utils.SparkTestPipelineOptions;
+import org.apache.beam.runners.spark.PipelineRule;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.io.AvroIO;
 import org.apache.beam.sdk.values.PCollection;
@@ -54,7 +54,7 @@ public class AvroPipelineTest {
   public final TemporaryFolder tmpDir = new TemporaryFolder();
 
   @Rule
-  public final SparkTestPipelineOptions pipelineOptions = new SparkTestPipelineOptions();
+  public final PipelineRule pipelineRule = PipelineRule.batch();
 
   @Before
   public void setUp() throws IOException {
@@ -72,7 +72,7 @@ public class AvroPipelineTest {
     savedRecord.put("siblingnames", Lists.newArrayList("Jimmy", "Jane"));
     populateGenericFile(Lists.newArrayList(savedRecord), schema);
 
-    Pipeline p = Pipeline.create(pipelineOptions.getOptions());
+    Pipeline p = pipelineRule.createPipeline();
     PCollection<GenericRecord> input = p.apply(
         AvroIO.Read.from(inputFile.getAbsolutePath()).withSchema(schema));
     input.apply(AvroIO.Write.to(outputDir.getAbsolutePath()).withSchema(schema));

http://git-wip-us.apache.org/repos/asf/beam/blob/4ca56806/runners/spark/src/test/java/org/apache/beam/runners/spark/io/NumShardsTest.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/io/NumShardsTest.java
b/runners/spark/src/test/java/org/apache/beam/runners/spark/io/NumShardsTest.java
index 34d6818..c936ed3 100644
--- a/runners/spark/src/test/java/org/apache/beam/runners/spark/io/NumShardsTest.java
+++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/io/NumShardsTest.java
@@ -30,8 +30,8 @@ import java.io.IOException;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Set;
+import org.apache.beam.runners.spark.PipelineRule;
 import org.apache.beam.runners.spark.examples.WordCount;
-import org.apache.beam.runners.spark.translation.streaming.utils.SparkTestPipelineOptions;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.io.TextIO;
@@ -59,7 +59,7 @@ public class NumShardsTest {
   public final TemporaryFolder tmpDir = new TemporaryFolder();
 
   @Rule
-  public final SparkTestPipelineOptions pipelineOptions = new SparkTestPipelineOptions();
+  public final PipelineRule pipelineRule = PipelineRule.batch();
 
   @Before
   public void setUp() throws IOException {
@@ -69,7 +69,7 @@ public class NumShardsTest {
 
   @Test
   public void testText() throws Exception {
-    Pipeline p = Pipeline.create(pipelineOptions.getOptions());
+    Pipeline p = pipelineRule.createPipeline();
     PCollection<String> inputWords = p.apply(Create.of(WORDS).withCoder(StringUtf8Coder.of()));
     PCollection<String> output = inputWords.apply(new WordCount.CountWords())
         .apply(MapElements.via(new WordCount.FormatAsTextFn()));

http://git-wip-us.apache.org/repos/asf/beam/blob/4ca56806/runners/spark/src/test/java/org/apache/beam/runners/spark/io/hadoop/HadoopFileFormatPipelineTest.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/io/hadoop/HadoopFileFormatPipelineTest.java
b/runners/spark/src/test/java/org/apache/beam/runners/spark/io/hadoop/HadoopFileFormatPipelineTest.java
index 9efc670..a5072d6 100644
--- a/runners/spark/src/test/java/org/apache/beam/runners/spark/io/hadoop/HadoopFileFormatPipelineTest.java
+++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/io/hadoop/HadoopFileFormatPipelineTest.java
@@ -22,8 +22,8 @@ import static org.junit.Assert.assertEquals;
 
 import java.io.File;
 import java.io.IOException;
+import org.apache.beam.runners.spark.PipelineRule;
 import org.apache.beam.runners.spark.coders.WritableCoder;
-import org.apache.beam.runners.spark.translation.streaming.utils.SparkTestPipelineOptions;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.values.KV;
@@ -52,7 +52,7 @@ public class HadoopFileFormatPipelineTest {
   private File outputFile;
 
   @Rule
-  public final SparkTestPipelineOptions pipelineOptions = new SparkTestPipelineOptions();
+  public final PipelineRule pipelineRule = PipelineRule.batch();
 
   @Rule
   public final TemporaryFolder tmpDir = new TemporaryFolder();
@@ -68,7 +68,7 @@ public class HadoopFileFormatPipelineTest {
   public void testSequenceFile() throws Exception {
     populateFile();
 
-    Pipeline p = Pipeline.create(pipelineOptions.getOptions());
+    Pipeline p = pipelineRule.createPipeline();
     @SuppressWarnings("unchecked")
     Class<? extends FileInputFormat<IntWritable, Text>> inputFormatClass =
         (Class<? extends FileInputFormat<IntWritable, Text>>)

http://git-wip-us.apache.org/repos/asf/beam/blob/4ca56806/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/StorageLevelTest.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/StorageLevelTest.java
b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/StorageLevelTest.java
index 48105e1..4dc5dee 100644
--- a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/StorageLevelTest.java
+++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/StorageLevelTest.java
@@ -17,7 +17,7 @@
  */
 package org.apache.beam.runners.spark.translation;
 
-import org.apache.beam.runners.spark.translation.streaming.utils.SparkTestPipelineOptions;
+import org.apache.beam.runners.spark.PipelineRule;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.transforms.Count;
@@ -32,12 +32,12 @@ import org.junit.Test;
 public class StorageLevelTest {
 
   @Rule
-  public final transient SparkTestPipelineOptions pipelineOptions = new SparkTestPipelineOptions();
+  public final transient PipelineRule pipelineRule = PipelineRule.batch();
 
   @Test
   public void test() throws Exception {
-    pipelineOptions.getOptions().setStorageLevel("DISK_ONLY");
-    Pipeline p = Pipeline.create(pipelineOptions.getOptions());
+    pipelineRule.getOptions().setStorageLevel("DISK_ONLY");
+    Pipeline p = pipelineRule.createPipeline();
 
     PCollection<String> pCollection = p.apply(Create.of("foo"));
 

http://git-wip-us.apache.org/repos/asf/beam/blob/4ca56806/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/CreateStreamTest.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/CreateStreamTest.java
b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/CreateStreamTest.java
index 9ee5cc5..f2783a1 100644
--- a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/CreateStreamTest.java
+++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/CreateStreamTest.java
@@ -24,10 +24,9 @@ import static org.junit.Assert.assertThat;
 
 import java.io.IOException;
 import java.io.Serializable;
+import org.apache.beam.runners.spark.PipelineRule;
 import org.apache.beam.runners.spark.ReuseSparkContextRule;
-import org.apache.beam.runners.spark.SparkPipelineOptions;
 import org.apache.beam.runners.spark.io.CreateStream;
-import org.apache.beam.runners.spark.translation.streaming.utils.SparkTestPipelineOptionsForStreaming;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.coders.VarIntCoder;
@@ -59,8 +58,6 @@ import org.joda.time.Instant;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
-import org.junit.rules.TemporaryFolder;
-import org.junit.rules.TestName;
 
 
 /**
@@ -76,27 +73,18 @@ import org.junit.rules.TestName;
 public class CreateStreamTest implements Serializable {
 
   @Rule
-  public transient TemporaryFolder checkpointParentDir = new TemporaryFolder();
+  public final transient PipelineRule pipelineRule = PipelineRule.streaming();
   @Rule
-  public transient SparkTestPipelineOptionsForStreaming commonOptions =
-      new SparkTestPipelineOptionsForStreaming();
+  public final transient ReuseSparkContextRule noContextResue = ReuseSparkContextRule.no();
   @Rule
-  public transient ReuseSparkContextRule noContextResue = ReuseSparkContextRule.no();
-  @Rule
-  public transient TestName testName = new TestName();
-  @Rule
-  public transient ExpectedException thrown = ExpectedException.none();
+  public final transient ExpectedException thrown = ExpectedException.none();
 
   @Test
   public void testLateDataAccumulating() throws IOException {
-    SparkPipelineOptions options = commonOptions.withTmpCheckpointDir(checkpointParentDir);
-    Pipeline p = Pipeline.create(options);
-    options.setJobName(testName.getMethodName());
-    Duration batchDuration = Duration.millis(options.getBatchIntervalMillis());
-
+    Pipeline p = pipelineRule.createPipeline();
     Instant instant = new Instant(0);
     CreateStream<TimestampedValue<Integer>> source =
-        CreateStream.<TimestampedValue<Integer>>withBatchInterval(batchDuration)
+        CreateStream.<TimestampedValue<Integer>>withBatchInterval(pipelineRule.batchDuration())
             .nextBatch()
             .advanceWatermarkForNextBatch(instant.plus(Duration.standardMinutes(6)))
             .nextBatch(
@@ -167,13 +155,9 @@ public class CreateStreamTest implements Serializable {
 
   @Test
   public void testDiscardingMode() throws IOException {
-    SparkPipelineOptions options = commonOptions.withTmpCheckpointDir(checkpointParentDir);
-    Pipeline p = Pipeline.create(options);
-    options.setJobName(testName.getMethodName());
-    Duration batchDuration = Duration.millis(options.getBatchIntervalMillis());
-
+    Pipeline p = pipelineRule.createPipeline();
     CreateStream<TimestampedValue<String>> source =
-        CreateStream.<TimestampedValue<String>>withBatchInterval(batchDuration)
+        CreateStream.<TimestampedValue<String>>withBatchInterval(pipelineRule.batchDuration())
             .nextBatch(
                 TimestampedValue.of("firstPane", new Instant(100)),
                 TimestampedValue.of("alsoFirstPane", new Instant(200)))
@@ -221,14 +205,10 @@ public class CreateStreamTest implements Serializable {
 
   @Test
   public void testFirstElementLate() throws IOException {
-    SparkPipelineOptions options = commonOptions.withTmpCheckpointDir(checkpointParentDir);
-    Pipeline p = Pipeline.create(options);
-    options.setJobName(testName.getMethodName());
-    Duration batchDuration = Duration.millis(options.getBatchIntervalMillis());
-
+    Pipeline p = pipelineRule.createPipeline();
     Instant lateElementTimestamp = new Instant(-1_000_000);
     CreateStream<TimestampedValue<String>> source =
-        CreateStream.<TimestampedValue<String>>withBatchInterval(batchDuration)
+        CreateStream.<TimestampedValue<String>>withBatchInterval(pipelineRule.batchDuration())
             .nextBatch()
             .advanceWatermarkForNextBatch(new Instant(0))
             .nextBatch(
@@ -261,14 +241,10 @@ public class CreateStreamTest implements Serializable {
 
   @Test
   public void testElementsAtAlmostPositiveInfinity() throws IOException {
-    SparkPipelineOptions options = commonOptions.withTmpCheckpointDir(checkpointParentDir);
-    Pipeline p = Pipeline.create(options);
-    options.setJobName(testName.getMethodName());
-    Duration batchDuration = Duration.millis(options.getBatchIntervalMillis());
-
+    Pipeline p = pipelineRule.createPipeline();
     Instant endOfGlobalWindow = GlobalWindow.INSTANCE.maxTimestamp();
     CreateStream<TimestampedValue<String>> source =
-        CreateStream.<TimestampedValue<String>>withBatchInterval(batchDuration)
+        CreateStream.<TimestampedValue<String>>withBatchInterval(pipelineRule.batchDuration())
             .nextBatch(
                 TimestampedValue.of("foo", endOfGlobalWindow),
                 TimestampedValue.of("bar", endOfGlobalWindow))
@@ -292,17 +268,13 @@ public class CreateStreamTest implements Serializable {
 
   @Test
   public void testMultipleStreams() throws IOException {
-    SparkPipelineOptions options = commonOptions.withTmpCheckpointDir(checkpointParentDir);
-    Pipeline p = Pipeline.create(options);
-    options.setJobName(testName.getMethodName());
-    Duration batchDuration = Duration.millis(options.getBatchIntervalMillis());
-
+    Pipeline p = pipelineRule.createPipeline();
     CreateStream<String> source =
-        CreateStream.<String>withBatchInterval(batchDuration)
+        CreateStream.<String>withBatchInterval(pipelineRule.batchDuration())
             .nextBatch("foo", "bar")
             .advanceNextBatchWatermarkToInfinity();
     CreateStream<Integer> other =
-        CreateStream.<Integer>withBatchInterval(batchDuration)
+        CreateStream.<Integer>withBatchInterval(pipelineRule.batchDuration())
             .nextBatch(1, 2, 3, 4)
             .advanceNextBatchWatermarkToInfinity();
 
@@ -327,14 +299,10 @@ public class CreateStreamTest implements Serializable {
 
   @Test
   public void testFlattenedWithWatermarkHold() throws IOException {
-    SparkPipelineOptions options = commonOptions.withTmpCheckpointDir(checkpointParentDir);
-    Pipeline p = Pipeline.create(options);
-    options.setJobName(testName.getMethodName());
-    Duration batchDuration = Duration.millis(options.getBatchIntervalMillis());
-
+    Pipeline p = pipelineRule.createPipeline();
     Instant instant = new Instant(0);
     CreateStream<TimestampedValue<Integer>> source1 =
-        CreateStream.<TimestampedValue<Integer>>withBatchInterval(batchDuration)
+        CreateStream.<TimestampedValue<Integer>>withBatchInterval(pipelineRule.batchDuration())
             .nextBatch()
             .advanceWatermarkForNextBatch(instant.plus(Duration.standardMinutes(5)))
             .nextBatch(
@@ -343,7 +311,7 @@ public class CreateStreamTest implements Serializable {
                 TimestampedValue.of(3, instant))
             .advanceWatermarkForNextBatch(instant.plus(Duration.standardMinutes(10)));
     CreateStream<TimestampedValue<Integer>> source2 =
-        CreateStream.<TimestampedValue<Integer>>withBatchInterval(batchDuration)
+        CreateStream.<TimestampedValue<Integer>>withBatchInterval(pipelineRule.batchDuration())
             .nextBatch()
             .advanceWatermarkForNextBatch(instant.plus(Duration.standardMinutes(1)))
             .nextBatch(
@@ -384,9 +352,8 @@ public class CreateStreamTest implements Serializable {
 
   @Test
   public void testElementAtPositiveInfinityThrows() {
-    Duration batchDuration = Duration.millis(commonOptions.getOptions().getBatchIntervalMillis());
     CreateStream<TimestampedValue<Integer>> source =
-        CreateStream.<TimestampedValue<Integer>>withBatchInterval(batchDuration)
+        CreateStream.<TimestampedValue<Integer>>withBatchInterval(pipelineRule.batchDuration())
             .nextBatch(TimestampedValue.of(-1, BoundedWindow.TIMESTAMP_MAX_VALUE.minus(1L)));
     thrown.expect(IllegalArgumentException.class);
     source.nextBatch(TimestampedValue.of(1, BoundedWindow.TIMESTAMP_MAX_VALUE));
@@ -394,9 +361,8 @@ public class CreateStreamTest implements Serializable {
 
   @Test
   public void testAdvanceWatermarkNonMonotonicThrows() {
-    Duration batchDuration = Duration.millis(commonOptions.getOptions().getBatchIntervalMillis());
     CreateStream<Integer> source =
-        CreateStream.<Integer>withBatchInterval(batchDuration)
+        CreateStream.<Integer>withBatchInterval(pipelineRule.batchDuration())
             .advanceWatermarkForNextBatch(new Instant(0L));
     thrown.expect(IllegalArgumentException.class);
     source.advanceWatermarkForNextBatch(new Instant(-1L));
@@ -404,9 +370,8 @@ public class CreateStreamTest implements Serializable {
 
   @Test
   public void testAdvanceWatermarkEqualToPositiveInfinityThrows() {
-    Duration batchDuration = Duration.millis(commonOptions.getOptions().getBatchIntervalMillis());
     CreateStream<Integer> source =
-        CreateStream.<Integer>withBatchInterval(batchDuration)
+        CreateStream.<Integer>withBatchInterval(pipelineRule.batchDuration())
             .advanceWatermarkForNextBatch(BoundedWindow.TIMESTAMP_MAX_VALUE.minus(1L));
     thrown.expect(IllegalArgumentException.class);
     source.advanceWatermarkForNextBatch(BoundedWindow.TIMESTAMP_MAX_VALUE);

http://git-wip-us.apache.org/repos/asf/beam/blob/4ca56806/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/utils/SparkTestPipelineOptions.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/utils/SparkTestPipelineOptions.java
b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/utils/SparkTestPipelineOptions.java
deleted file mode 100644
index efc17d3..0000000
--- a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/utils/SparkTestPipelineOptions.java
+++ /dev/null
@@ -1,42 +0,0 @@
-/*
- * 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.runners.spark.translation.streaming.utils;
-
-import org.apache.beam.runners.spark.SparkPipelineOptions;
-import org.apache.beam.runners.spark.TestSparkRunner;
-import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.junit.rules.ExternalResource;
-
-/**
- * A rule to create a common {@link SparkPipelineOptions} test options for spark-runner.
- */
-public class SparkTestPipelineOptions extends ExternalResource {
-
-  protected final SparkPipelineOptions options =
-      PipelineOptionsFactory.as(SparkPipelineOptions.class);
-
-  @Override
-  protected void before() throws Throwable {
-    options.setRunner(TestSparkRunner.class);
-    options.setEnableSparkMetricSinks(false);
-  }
-
-  public SparkPipelineOptions getOptions() {
-    return options;
-  }
-}

http://git-wip-us.apache.org/repos/asf/beam/blob/4ca56806/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/utils/SparkTestPipelineOptionsForStreaming.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/utils/SparkTestPipelineOptionsForStreaming.java
b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/utils/SparkTestPipelineOptionsForStreaming.java
deleted file mode 100644
index dd3e4c8..0000000
--- a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/utils/SparkTestPipelineOptionsForStreaming.java
+++ /dev/null
@@ -1,48 +0,0 @@
-/*
- * 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.runners.spark.translation.streaming.utils;
-
-import java.io.IOException;
-import org.apache.beam.runners.spark.SparkPipelineOptions;
-import org.joda.time.Duration;
-import org.junit.rules.TemporaryFolder;
-
-
-/**
- * A rule to create a common {@link SparkPipelineOptions} for testing streaming pipelines.
- */
-public class SparkTestPipelineOptionsForStreaming extends SparkTestPipelineOptions {
-
-  private static final int DEFAULT_NUMBER_OF_BATCHES_TIMEOUT = 5;
-
-  public SparkPipelineOptions withTmpCheckpointDir(TemporaryFolder parent)
-      throws IOException {
-    // tests use JUnit's TemporaryFolder path in the form of: /.../junit/...
-    options.setCheckpointDir(parent.newFolder(options.getJobName()).toURI().toURL().toString());
-    options.setForceStreaming(true);
-    // set the default timeout to DEFAULT_NUMBER_OF_BATCHES_TIMEOUT x batchDuration
-    // to allow pipelines to finish.
-    Duration batchDuration = Duration.millis(options.getBatchIntervalMillis());
-    // set the checkpoint duration to match interval.
-    options.setCheckpointDurationMillis(batchDuration.getMillis());
-    long forcedTimeout = batchDuration.multipliedBy(DEFAULT_NUMBER_OF_BATCHES_TIMEOUT).getMillis();
-    options.setForcedTimeout(forcedTimeout);
-    return options;
-  }
-}


Mime
View raw message