It is recommended to tag hand-selected tests for this purpose using the {@link + * ValidatesRunner} {@link Category} annotation, as each test run against a pipeline runner will + * utilize resources of that pipeline runner. + * + *
In order to run tests on a pipeline runner, the following conditions must be met: + * + *
{@code [
+ * "--runner=TestDataflowRunner",
+ * "--project=mygcpproject",
+ * "--stagingLocation=gs://mygcsbucket/path"
+ * ]}
+ * Note that the set of pipeline options required is pipeline runner specific.
+ * Use {@link PAssert} for tests, as it integrates with this test harness in both direct and + * remote execution modes. For example: + * + *
+ * {@literal @Rule}
+ * public final transient TestPipeline p = TestPipeline.create();
+ *
+ * {@literal @Test}
+ * {@literal @Category}(NeedsRunner.class)
+ * public void myPipelineTest() throws Exception {
+ * final PCollection<String> pCollection = pipeline.apply(...)
+ * PAssert.that(pCollection).containsInAnyOrder(...);
+ * pipeline.run();
+ * }
+ * For pipeline runners, it is required that they must throw an {@link AssertionError} containing + * the message from the {@link PAssert} that failed. + * + *
See also the Testing documentation
+ * section.
+ */
+public class TestPipelineExtension extends Pipeline implements TestRule {
+
+ private final PipelineOptions options;
+
+ private static class PipelineRunEnforcement {
+
+ @SuppressWarnings("WeakerAccess")
+ protected boolean enableAutoRunIfMissing;
+
+ protected final Pipeline pipeline;
+
+ protected boolean runAttempted;
+
+ private PipelineRunEnforcement(final Pipeline pipeline) {
+ this.pipeline = pipeline;
+ }
+
+ protected void enableAutoRunIfMissing(final boolean enable) {
+ enableAutoRunIfMissing = enable;
+ }
+
+ protected void beforePipelineExecution() {
+ runAttempted = true;
+ }
+
+ protected void afterPipelineExecution() {}
+
+ protected void afterUserCodeFinished() {
+ if (!runAttempted && enableAutoRunIfMissing) {
+ pipeline.run().waitUntilFinish();
+ }
+ }
+ }
+
+ private static class PipelineAbandonedNodeEnforcement extends PipelineRunEnforcement {
+
+ // Null until the pipeline has been run
+ @Nullable private List Use {@link PAssert} to add tests, then call {@link Pipeline#run} to execute the pipeline and
+ * check the tests.
+ */
+ public static TestPipelineExtension create() {
+ return fromOptions(testingPipelineOptions());
+ }
+
+ public static TestPipelineExtension fromOptions(PipelineOptions options) {
+ return new TestPipelineExtension(options);
+ }
+
+ private TestPipelineExtension(final PipelineOptions options) {
+ super(options);
+ this.options = options;
+ }
+
+ @Override
+ public PipelineOptions getOptions() {
+ return this.options;
+ }
+
+ @Override
+ public Statement apply(final Statement statement, final Description description) {
+ return new Statement() {
+
+ private void setDeducedEnforcementLevel() {
+ // if the enforcement level has not been set by the user do auto-inference
+ if (!enforcement.isPresent()) {
+
+ final boolean annotatedWithNeedsRunner =
+ FluentIterable.from(description.getAnnotations())
+ .filter(Annotations.Predicates.isAnnotationOfType(Category.class))
+ .anyMatch(Annotations.Predicates.isCategoryOf(NeedsRunner.class, true));
+
+ final boolean crashingRunner = CrashingRunner.class.isAssignableFrom(options.getRunner());
+
+ checkState(
+ !(annotatedWithNeedsRunner && crashingRunner),
+ "The test was annotated with a [@%s] / [@%s] while the runner "
+ + "was set to [%s]. Please re-check your configuration.",
+ NeedsRunner.class.getSimpleName(),
+ ValidatesRunner.class.getSimpleName(),
+ CrashingRunner.class.getSimpleName());
+
+ enableAbandonedNodeEnforcement(annotatedWithNeedsRunner || !crashingRunner);
+ }
+ }
+
+ @Override
+ public void evaluate() throws Throwable {
+ options.as(ApplicationNameOptions.class).setAppName(getAppName(description));
+
+ setDeducedEnforcementLevel();
+
+ // statement.evaluate() essentially runs the user code contained in the unit test at hand.
+ // Exceptions thrown during the execution of the user's test code will propagate here,
+ // unless the user explicitly handles them with a "catch" clause in his code. If the
+ // exception is handled by a user's "catch" clause, is does not interrupt the flow and
+ // we move on to invoking the configured enforcements.
+ // If the user does not handle a thrown exception, it will propagate here and interrupt
+ // the flow, preventing the enforcement(s) from being activated.
+ // The motivation for this is avoiding enforcements over faulty pipelines.
+ statement.evaluate();
+ enforcement.get().afterUserCodeFinished();
+ }
+ };
+ }
+
+ /**
+ * Runs this {@link TestPipelineExtension}, unwrapping any {@code AssertionError} that is raised during
+ * testing.
+ */
+ @Override
+ public PipelineResult run() {
+ return run(getOptions());
+ }
+
+ /** Like {@link #run} but with the given potentially modified options. */
+ @Override
+ public PipelineResult run(PipelineOptions options) {
+ checkState(
+ enforcement.isPresent(),
+ "Is your TestPipeline declaration missing a @Rule annotation? Usage: "
+ + "@Rule public final transient TestPipeline pipeline = TestPipeline.create();");
+
+ final PipelineResult pipelineResult;
+ try {
+ enforcement.get().beforePipelineExecution();
+ PipelineOptions updatedOptions =
+ MAPPER.convertValue(MAPPER.valueToTree(options), PipelineOptions.class);
+ updatedOptions
+ .as(TestValueProviderOptions.class)
+ .setProviderRuntimeValues(StaticValueProvider.of(providerRuntimeValues));
+ pipelineResult = super.run(updatedOptions);
+ verifyPAssertsSucceeded(this, pipelineResult);
+ } catch (RuntimeException exc) {
+ Throwable cause = exc.getCause();
+ if (cause instanceof AssertionError) {
+ throw (AssertionError) cause;
+ } else {
+ throw exc;
+ }
+ }
+
+ // If we reach this point, the pipeline has been run and no exceptions have been thrown during
+ // its execution.
+ enforcement.get().afterPipelineExecution();
+ return pipelineResult;
+ }
+
+ /** Implementation detail of {@link #newProvider}, do not use. */
+ @Internal
+ public interface TestValueProviderOptions extends PipelineOptions {
+ ValueProvider Note this only runs for runners which support Metrics. Runners which do not should verify
+ * this in some other way. See: https://issues.apache.org/jira/browse/BEAM-2001
+ */
+ public static void verifyPAssertsSucceeded(Pipeline pipeline, PipelineResult pipelineResult) {
+ if (MetricsEnvironment.isMetricsSupported()) {
+ long expectedNumberOfAssertions = (long) PAssert.countAsserts(pipeline);
+
+ long successfulAssertions = 0;
+ Iterable NOTE: This was copied from the Apache Beam project from a separate file only for visibility
+ * reasons (it's package-private there).
+ */
+ static class Annotations {
+
+ /** Annotation predicates. */
+ static class Predicates {
+
+ static Predicate This class implements {@link Serializable} so that test {@link DoFn} classes may be inlined.
*/
-// TODO(weiminyu): Upgrade to JUnit5 when TestPipeline is upgraded. It is also easy to adapt with
-// a wrapper.
@RunWith(JUnit4.class)
public class ExportloadingTransformsTest implements Serializable {
private static final DateTime START_TIME = DateTime.parse("2000-01-01T00:00:00.0Z");
@@ -72,8 +70,8 @@ public class ExportloadingTransformsTest implements Serializable {
@Rule public final transient InjectRule injectRule = new InjectRule();
@Rule
- public final transient TestPipeline pipeline =
- TestPipeline.create().enableAbandonedNodeEnforcement(true);
+ public final transient TestPipelineExtension pipeline =
+ TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
private FakeClock fakeClock;
private transient BackupTestStore store;
diff --git a/core/src/test/java/google/registry/beam/initsql/LoadDatastoreSnapshotTest.java b/core/src/test/java/google/registry/beam/initsql/LoadDatastoreSnapshotTest.java
index 8399d8dcd..477160adf 100644
--- a/core/src/test/java/google/registry/beam/initsql/LoadDatastoreSnapshotTest.java
+++ b/core/src/test/java/google/registry/beam/initsql/LoadDatastoreSnapshotTest.java
@@ -22,6 +22,7 @@ import com.google.appengine.api.datastore.Entity;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.googlecode.objectify.Key;
+import google.registry.beam.TestPipelineExtension;
import google.registry.model.contact.ContactResource;
import google.registry.model.domain.DomainAuthInfo;
import google.registry.model.domain.DomainBase;
@@ -32,7 +33,6 @@ import google.registry.testing.FakeClock;
import google.registry.testing.InjectRule;
import java.io.File;
import org.apache.beam.sdk.testing.NeedsRunner;
-import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.joda.time.DateTime;
@@ -85,8 +85,8 @@ public class LoadDatastoreSnapshotTest {
@Rule public final transient InjectRule injectRule = new InjectRule();
@Rule
- public final transient TestPipeline pipeline =
- TestPipeline.create().enableAbandonedNodeEnforcement(true);
+ public final transient TestPipelineExtension pipeline =
+ TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
private FakeClock fakeClock;
private File exportRootDir;
diff --git a/core/src/test/java/google/registry/beam/initsql/WriteToSqlTest.java b/core/src/test/java/google/registry/beam/initsql/WriteToSqlTest.java
index 95ce904a8..a1821dd7e 100644
--- a/core/src/test/java/google/registry/beam/initsql/WriteToSqlTest.java
+++ b/core/src/test/java/google/registry/beam/initsql/WriteToSqlTest.java
@@ -20,6 +20,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
import com.google.appengine.api.datastore.Entity;
import com.google.common.collect.ImmutableList;
import google.registry.backup.VersionedEntity;
+import google.registry.beam.TestPipelineExtension;
import google.registry.model.contact.ContactResource;
import google.registry.model.ofy.Ofy;
import google.registry.model.registrar.Registrar;
@@ -35,7 +36,6 @@ import java.io.PrintStream;
import java.io.Serializable;
import java.util.stream.Collectors;
import org.apache.beam.sdk.testing.NeedsRunner;
-import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Create;
import org.joda.time.DateTime;
import org.junit.Before;
@@ -65,8 +65,8 @@ public class WriteToSqlTest implements Serializable {
@Rule public transient TemporaryFolder temporaryFolder = new TemporaryFolder();
@Rule
- public final transient TestPipeline pipeline =
- TestPipeline.create().enableAbandonedNodeEnforcement(true);
+ public final transient TestPipelineExtension pipeline =
+ TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
private ImmutableListPTransforms,
+ * PAsserts included, that were not executed by the pipeline runner. Abandoned nodes are
+ * most likely to occur due to the one of the following scenarios:
+ *
+ *
+ *
+ *
+ * Abandoned node detection is automatically enabled when a real pipeline runner (i.e. not a
+ * {@link CrashingRunner}) and/or a {@link NeedsRunner} or a {@link ValidatesRunner} annotation
+ * are detected.
+ */
+ public TestPipelineExtension enableAbandonedNodeEnforcement(final boolean enable) {
+ enforcement =
+ enable
+ ? Optional.of(new PipelineAbandonedNodeEnforcement(this))
+ : Optional.of(new PipelineRunEnforcement(this));
+
+ return this;
+ }
+
+ /**
+ * If enabled, a pipeline.run() statement at the end of a test.
+ * pipeline.run() statement will be added automatically in case it is
+ * missing in the test.
+ */
+ public TestPipelineExtension enableAutoRunIfMissing(final boolean enable) {
+ enforcement.get().enableAutoRunIfMissing(enable);
+ return this;
+ }
+
+ @Override
+ public String toString() {
+ return "TestPipeline#" + options.as(ApplicationNameOptions.class).getAppName();
+ }
+
+ /** Creates {@link PipelineOptions} for testing. */
+ public static PipelineOptions testingPipelineOptions() {
+ try {
+ @Nullable
+ String beamTestPipelineOptions = System.getProperty(PROPERTY_BEAM_TEST_PIPELINE_OPTIONS);
+
+ PipelineOptions options =
+ Strings.isNullOrEmpty(beamTestPipelineOptions)
+ ? PipelineOptionsFactory.create()
+ : PipelineOptionsFactory.fromArgs(
+ MAPPER.readValue(beamTestPipelineOptions, String[].class))
+ .as(TestPipelineOptions.class);
+
+ // If no options were specified, set some reasonable defaults
+ if (Strings.isNullOrEmpty(beamTestPipelineOptions)) {
+ // If there are no provided options, check to see if a dummy runner should be used.
+ String useDefaultDummy = System.getProperty(PROPERTY_USE_DEFAULT_DUMMY_RUNNER);
+ if (!Strings.isNullOrEmpty(useDefaultDummy) && Boolean.valueOf(useDefaultDummy)) {
+ options.setRunner(CrashingRunner.class);
+ }
+ }
+ options.setStableUniqueNames(CheckEnabled.ERROR);
+
+ FileSystems.setDefaultPipelineOptions(options);
+ return options;
+ } catch (IOException e) {
+ throw new RuntimeException(
+ "Unable to instantiate test options from system property "
+ + PROPERTY_BEAM_TEST_PIPELINE_OPTIONS
+ + ":"
+ + System.getProperty(PROPERTY_BEAM_TEST_PIPELINE_OPTIONS),
+ e);
+ }
+ }
+
+ /** Returns the class + method name of the test. */
+ private String getAppName(Description description) {
+ String methodName = description.getMethodName();
+ Class testClass = description.getTestClass();
+ if (testClass.isMemberClass()) {
+ return String.format(
+ "%s$%s-%s",
+ testClass.getEnclosingClass().getSimpleName(), testClass.getSimpleName(), methodName);
+ } else {
+ return String.format("%s-%s", testClass.getSimpleName(), methodName);
+ }
+ }
+
+ /**
+ * Verifies all {{@link PAssert PAsserts}} in the pipeline have been executed and were successful.
+ *
+ *