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

The following examples show how to use org.apache.beam.sdk.testing.TestPipeline#fromOptions() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source File: BigQueryIOReadTest.java    From beam with Apache License 2.0 6 votes vote down vote up 
@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);
}
Example 2
Source File: BigQueryIOStorageReadTest.java    From beam with Apache License 2.0 6 votes vote down vote up 
@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);
}
Example 3
Source File: BigQueryIOWriteTest.java    From beam with Apache License 2.0 6 votes vote down vote up 
@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);
}
Example 4
Source File: BigQueryIOStorageQueryTest.java    From beam with Apache License 2.0 6 votes vote down vote up 
@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);
}
Example 5
Source File: PubsubIOTest.java    From beam with Apache License 2.0 6 votes vote down vote up 
@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;
}
Example 6
Source File: FlinkTransformOverridesTest.java    From beam with Apache License 2.0 6 votes vote down vote up 
@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));
}
Example 7
Source File: DataflowRunnerTest.java    From beam with Apache License 2.0 6 votes vote down vote up 
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());
}
Example 8
Source File: TestUtils.java    From beam with Apache License 2.0 5 votes vote down vote up 
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;
}
Example 9
Source File: DataflowRunnerTest.java    From beam with Apache License 2.0 5 votes vote down vote up 
@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());
}