Java Code Examples for org.apache.beam.sdk.testing.TestPipeline#fromOptions()

@Override
public Statement apply(final Statement base, final Description description) {
  // We need to set up the temporary folder, and then set up the TestPipeline based on the
  // chosen folder. Unfortunately, since rule evaluation order is unspecified and unrelated
  // to field order, and is separate from construction, that requires manually creating this
  // TestRule.
  Statement withPipeline =
      new Statement() {
        @Override
        public void evaluate() throws Throwable {
          options = TestPipeline.testingPipelineOptions();
          options.as(BigQueryOptions.class).setProject("project-id");
          options
              .as(BigQueryOptions.class)
              .setTempLocation(testFolder.getRoot().getAbsolutePath());
          p = TestPipeline.fromOptions(options);
          p.apply(base, description).evaluate();
        }
      };
  return testFolder.apply(withPipeline, description);
}
 

@Override
public Statement apply(Statement base, Description description) {
  // We need to set up the temporary folder, and then set up the TestPipeline based on the
  // chosen folder. Unfortunately, since rule evaluation order is unspecified and unrelated
  // to field order, and is separate from construction, that requires manually creating this
  // TestRule.
  Statement withPipeline =
      new Statement() {
        @Override
        public void evaluate() throws Throwable {
          options = TestPipeline.testingPipelineOptions();
          options.as(BigQueryOptions.class).setProject("project-id");
          options
              .as(BigQueryOptions.class)
              .setTempLocation(testFolder.getRoot().getAbsolutePath());
          p = TestPipeline.fromOptions(options);
          p.apply(base, description).evaluate();
        }
      };
  return testFolder.apply(withPipeline, description);
}
 

@Override
public Statement apply(final Statement base, final Description description) {
  // We need to set up the temporary folder, and then set up the TestPipeline based on the
  // chosen folder. Unfortunately, since rule evaluation order is unspecified and unrelated
  // to field order, and is separate from construction, that requires manually creating this
  // TestRule.
  Statement withPipeline =
      new Statement() {
        @Override
        public void evaluate() throws Throwable {
          options = TestPipeline.testingPipelineOptions();
          options.as(BigQueryOptions.class).setProject("project-id");
          options
              .as(BigQueryOptions.class)
              .setTempLocation(testFolder.getRoot().getAbsolutePath());
          p = TestPipeline.fromOptions(options);
          p.apply(base, description).evaluate();
        }
      };
  return testFolder.apply(withPipeline, description);
}
 

@Override
public Statement apply(Statement base, Description description) {
  // We need to set up the temporary folder, and then set up the TestPipeline based on the
  // chosen folder. Unfortunately, since rule evaluation order is unspecified and unrelated
  // to field order, and is separate from construction, that requires manually creating this
  // TestRule.
  Statement withPipeline =
      new Statement() {
        @Override
        public void evaluate() throws Throwable {
          options = TestPipeline.testingPipelineOptions().as(BigQueryOptions.class);
          options.setProject("project-id");
          options.setTempLocation(testFolder.getRoot().getAbsolutePath());
          p = TestPipeline.fromOptions(options);
          p.apply(base, description).evaluate();
        }
      };

  return testFolder.apply(withPipeline, description);
}
 

@Override
public Statement apply(final Statement base, final Description description) {
  // We need to set up the temporary folder, and then set up the TestPipeline based on the
  // chosen folder. Unfortunately, since rule evaluation order is unspecified and unrelated
  // to field order, and is separate from construction, that requires manually creating this
  // TestRule.
  Statement withPipeline =
      new Statement() {
        @Override
        public void evaluate() throws Throwable {
          options = TestPipeline.testingPipelineOptions();
          options.as(PubsubOptions.class).setProject("test-project");
          readPipeline = TestPipeline.fromOptions(options);
          readPipeline.apply(base, description).evaluate();
        }
      };
  return withPipeline;
}
 

@Test
public void testRunnerDeterminedSharding() {
  FlinkPipelineOptions options = PipelineOptionsFactory.as(FlinkPipelineOptions.class);
  options.setRunner(TestFlinkRunner.class);
  options.setFlinkMaster("[auto]");
  options.setParallelism(5);

  TestPipeline p = TestPipeline.fromOptions(options);

  StreamingShardedWriteFactory<Object, Void, Object> factory =
      new StreamingShardedWriteFactory<>(p.getOptions());

  WriteFiles<Object, Void, Object> original = WriteFiles.to(new TestSink(tmpFolder.toString()));
  @SuppressWarnings("unchecked")
  PCollection<Object> objs = (PCollection) p.apply(Create.empty(VoidCoder.of()));
  AppliedPTransform<PCollection<Object>, WriteFilesResult<Void>, WriteFiles<Object, Void, Object>>
      originalApplication =
          AppliedPTransform.of("writefiles", objs.expand(), Collections.emptyMap(), original, p);

  WriteFiles<Object, Void, Object> replacement =
      (WriteFiles<Object, Void, Object>)
          factory.getReplacementTransform(originalApplication).getTransform();

  assertThat(replacement, not(equalTo((Object) original)));
  assertThat(replacement.getNumShardsProvider().get(), is(10));
}
 

private void testStreamingWriteOverride(PipelineOptions options, int expectedNumShards) {
  TestPipeline p = TestPipeline.fromOptions(options);

  StreamingShardedWriteFactory<Object, Void, Object> factory =
      new StreamingShardedWriteFactory<>(p.getOptions());
  WriteFiles<Object, Void, Object> original = WriteFiles.to(new TestSink(tmpFolder.toString()));
  PCollection<Object> objs = (PCollection) p.apply(Create.empty(VoidCoder.of()));
  AppliedPTransform<PCollection<Object>, WriteFilesResult<Void>, WriteFiles<Object, Void, Object>>
      originalApplication =
          AppliedPTransform.of("writefiles", objs.expand(), Collections.emptyMap(), original, p);

  WriteFiles<Object, Void, Object> replacement =
      (WriteFiles<Object, Void, Object>)
          factory.getReplacementTransform(originalApplication).getTransform();
  assertThat(replacement, not(equalTo((Object) original)));
  assertThat(replacement.getNumShardsProvider().get(), equalTo(expectedNumShards));

  WriteFilesResult<Void> originalResult = objs.apply(original);
  WriteFilesResult<Void> replacementResult = objs.apply(replacement);
  Map<PValue, ReplacementOutput> res =
      factory.mapOutputs(originalResult.expand(), replacementResult);
  assertEquals(1, res.size());
  assertEquals(
      originalResult.getPerDestinationOutputFilenames(),
      res.get(replacementResult.getPerDestinationOutputFilenames()).getOriginal().getValue());
}
 

public static TestPipeline createTestPipeline() {
  final PipelineOptions pipelineOptions = PipelineOptionsFactory.create();
  pipelineOptions
      .as(EuphoriaOptions.class)
      .setTranslatorProvider(new PrimitiveOutputTranslatorProvider());
  final TestPipeline testPipeline = TestPipeline.fromOptions(pipelineOptions);
  testPipeline
      .getCoderRegistry()
      .registerCoderForClass(Object.class, KryoCoder.of(pipelineOptions));
  return testPipeline;
}
 

@Test
public void testPathValidatorOverride() {
  String[] args =
      new String[] {
        "--runner=DataflowRunner",
        "--region=some-region-1",
        "--tempLocation=/tmp/testing",
        "--project=test-project",
        "--credentialFactoryClass=" + NoopCredentialFactory.class.getName(),
        "--pathValidatorClass=" + NoopPathValidator.class.getName(),
      };
  // Should not crash, because gcpTempLocation should get set from tempLocation
  TestPipeline.fromOptions(PipelineOptionsFactory.fromArgs(args).create());
}