beam-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From amits...@apache.org
Subject [04/23] beam git commit: Use TestSparkRunner in tests.
Date Tue, 28 Feb 2017 22:35:11 GMT
Use TestSparkRunner in tests.


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

Branch: refs/heads/master
Commit: 2bcd40c206d6b615135375276b249364fedf32af
Parents: c25a02f
Author: Sela <ansela@paypal.com>
Authored: Sat Feb 18 22:10:50 2017 +0200
Committer: Sela <ansela@paypal.com>
Committed: Wed Mar 1 00:17:59 2017 +0200

----------------------------------------------------------------------
 .../beam/runners/spark/TestSparkRunner.java     | 92 ++++++++++++--------
 .../utils/SparkTestPipelineOptions.java         |  4 +-
 .../SparkTestPipelineOptionsForStreaming.java   | 11 +++
 3 files changed, 69 insertions(+), 38 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/beam/blob/2bcd40c2/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 8b8f9ba..24bc038 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
@@ -21,12 +21,12 @@ package org.apache.beam.runners.spark;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.is;
 
+import java.io.File;
+import java.io.IOException;
 import org.apache.beam.runners.core.UnboundedReadFromBoundedSource;
-import org.apache.beam.runners.spark.aggregators.AggregatorsAccumulator;
-import org.apache.beam.runners.spark.metrics.SparkMetricsContainer;
+import org.apache.beam.runners.spark.aggregators.AccumulatorSingleton;
 import org.apache.beam.runners.spark.util.GlobalWatermarkHolder;
 import org.apache.beam.sdk.Pipeline;
-import org.apache.beam.sdk.PipelineResult.State;
 import org.apache.beam.sdk.io.BoundedReadFromUnboundedSource;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.options.PipelineOptionsValidator;
@@ -40,6 +40,10 @@ import org.apache.beam.sdk.values.PBegin;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PInput;
 import org.apache.beam.sdk.values.POutput;
+import org.apache.commons.io.FileUtils;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 
 /**
@@ -65,6 +69,8 @@ import org.apache.beam.sdk.values.POutput;
  */
 public final class TestSparkRunner extends PipelineRunner<SparkPipelineResult> {
 
+  private static final Logger LOG = LoggerFactory.getLogger(TestSparkRunner.class);
+
   private SparkRunner delegate;
   private boolean isForceStreaming;
   private int expectedNumberOfAssertions = 0;
@@ -87,7 +93,7 @@ public final class TestSparkRunner extends PipelineRunner<SparkPipelineResult>
{
   @SuppressWarnings("unchecked")
   @Override
   public <OutputT extends POutput, InputT extends PInput> OutputT apply(
-          PTransform<InputT, OutputT> transform, InputT input) {
+      PTransform<InputT, OutputT> transform, InputT input) {
     // if the pipeline forces execution as a streaming pipeline,
     // and the source is an adapted unbounded source (as bounded),
     // read it as unbounded source via UnboundedReadFromBoundedSource.
@@ -108,38 +114,52 @@ public final class TestSparkRunner extends PipelineRunner<SparkPipelineResult>
{
 
   @Override
   public SparkPipelineResult run(Pipeline pipeline) {
-    // clear state of Aggregators, Metrics and Watermarks if exists.
-    AggregatorsAccumulator.clear();
-    SparkMetricsContainer.clear();
-    GlobalWatermarkHolder.clear();
-
-    TestPipelineOptions testPipelineOptions = pipeline.getOptions().as(TestPipelineOptions.class);
-    SparkPipelineResult result = delegate.run(pipeline);
-    result.waitUntilFinish();
-
-
-    // make sure the test pipeline finished successfully.
-    State resultState = result.getState();
-    assertThat(
-        String.format("Test pipeline result state was %s instead of %s", resultState, State.DONE),
-        resultState,
-        is(State.DONE));
-    assertThat(result, testPipelineOptions.getOnCreateMatcher());
-    assertThat(result, testPipelineOptions.getOnSuccessMatcher());
-
-    // if the pipeline was executed in streaming mode, validate aggregators.
-    if (isForceStreaming) {
-      // validate assertion succeeded (at least once).
-      int success = result.getAggregatorValue(PAssert.SUCCESS_COUNTER, Integer.class);
-      assertThat(
-          String.format(
-              "Expected %d successful assertions, but found %d.",
-              expectedNumberOfAssertions, success),
-          success,
-          is(expectedNumberOfAssertions));
-      // validate assertion didn't fail.
-      int failure = result.getAggregatorValue(PAssert.FAILURE_COUNTER, Integer.class);
-      assertThat("Failure aggregator should be zero.", failure, is(0));
+    SparkPipelineOptions sparkOptions = pipeline.getOptions().as(SparkPipelineOptions.class);
+    long timeout = sparkOptions.getForcedTimeout();
+    SparkPipelineResult result = null;
+    try {
+      // clear state of Accumulators and Aggregators.
+      AccumulatorSingleton.clear();
+      GlobalWatermarkHolder.clear();
+
+      TestPipelineOptions testPipelineOptions = pipeline.getOptions().as(TestPipelineOptions.class);
+      LOG.info("About to run test pipeline " + sparkOptions.getJobName());
+      result = delegate.run(pipeline);
+      result.waitUntilFinish(Duration.millis(timeout));
+
+      assertThat(result, testPipelineOptions.getOnCreateMatcher());
+      assertThat(result, testPipelineOptions.getOnSuccessMatcher());
+
+      // if the pipeline was executed in streaming mode, validate aggregators.
+      if (isForceStreaming) {
+        // validate assertion succeeded (at least once).
+        int successAssertions = result.getAggregatorValue(PAssert.SUCCESS_COUNTER, Integer.class);
+        assertThat(
+            String.format(
+                "Expected %d successful assertions, but found %d.",
+                expectedNumberOfAssertions, successAssertions),
+            successAssertions,
+            is(expectedNumberOfAssertions));
+        // validate assertion didn't fail.
+        int failedAssertions = result.getAggregatorValue(PAssert.FAILURE_COUNTER, Integer.class);
+        assertThat(
+            String.format("Found %d failed assertions.", failedAssertions),
+            failedAssertions,
+            is(0));
+
+        LOG.info(
+            String.format(
+                "Successfully asserted pipeline %s with %d successful assertions.",
+                sparkOptions.getJobName(),
+                successAssertions));
+      }
+    } finally {
+      try {
+        // cleanup checkpoint dir.
+        FileUtils.deleteDirectory(new File(sparkOptions.getCheckpointDir()));
+      } catch (IOException e) {
+        throw new RuntimeException("Failed to clear checkpoint tmp dir.", e);
+      }
     }
     return result;
   }

http://git-wip-us.apache.org/repos/asf/beam/blob/2bcd40c2/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
index 2da9888..efc17d3 100644
--- 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
@@ -18,7 +18,7 @@
 package org.apache.beam.runners.spark.translation.streaming.utils;
 
 import org.apache.beam.runners.spark.SparkPipelineOptions;
-import org.apache.beam.runners.spark.SparkRunner;
+import org.apache.beam.runners.spark.TestSparkRunner;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.junit.rules.ExternalResource;
 
@@ -32,7 +32,7 @@ public class SparkTestPipelineOptions extends ExternalResource {
 
   @Override
   protected void before() throws Throwable {
-    options.setRunner(SparkRunner.class);
+    options.setRunner(TestSparkRunner.class);
     options.setEnableSparkMetricSinks(false);
   }
 

http://git-wip-us.apache.org/repos/asf/beam/blob/2bcd40c2/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
index 28f6d5d..dd3e4c8 100644
--- 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
@@ -20,6 +20,7 @@ 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;
 
 
@@ -28,10 +29,20 @@ import org.junit.rules.TemporaryFolder;
  */
 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