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

private void exportAndImportDbAtTime(String sourceDb, String destDb,
                                     String jobIdName, String ts,
                                     TestPipeline exportPipeline,
                                     TestPipeline importPipeline) {
  ValueProvider.StaticValueProvider<String> destination = ValueProvider.StaticValueProvider
      .of(tmpDir);
  ValueProvider.StaticValueProvider<String> jobId = ValueProvider.StaticValueProvider
      .of(jobIdName);
  ValueProvider.StaticValueProvider<String> source = ValueProvider.StaticValueProvider
      .of(tmpDir + "/" + jobIdName);
  ValueProvider.StaticValueProvider<String> timestamp = ValueProvider.StaticValueProvider.of(ts);
  SpannerConfig sourceConfig = spannerServer.getSpannerConfig(sourceDb);
  exportPipeline.apply("Export", new ExportTransform(sourceConfig, destination,
                                                     jobId, timestamp));
  PipelineResult exportResult = exportPipeline.run();
  exportResult.waitUntilFinish();

  SpannerConfig copyConfig = spannerServer.getSpannerConfig(destDb);
  importPipeline.apply("Import", new ImportTransform(
      copyConfig, source, ValueProvider.StaticValueProvider.of(true),
      ValueProvider.StaticValueProvider.of(true),
      ValueProvider.StaticValueProvider.of(true)));
  PipelineResult importResult = importPipeline.run();
  importResult.waitUntilFinish();
}
 

private void compareDbs(String sourceDb, String destDb, TestPipeline comparePipeline) {
  SpannerConfig sourceConfig = spannerServer.getSpannerConfig(sourceDb);
  SpannerConfig copyConfig = spannerServer.getSpannerConfig(destDb);
  PCollection<Long> mismatchCount = comparePipeline
      .apply("Compare", new CompareDatabases(sourceConfig, copyConfig));
  PAssert.that(mismatchCount).satisfies((x) -> {
    assertEquals(Lists.newArrayList(x), Lists.newArrayList(0L));
    return null;
  });
  PipelineResult compareResult = comparePipeline.run();
  compareResult.waitUntilFinish();

  Ddl sourceDdl = readDdl(sourceDb);
  Ddl destinationDdl = readDdl(destDb);

  assertThat(sourceDdl.prettyPrint(), equalToIgnoringWhiteSpace(destinationDdl.prettyPrint()));
}
 

@Test
public void testEquals() {
  TestPipeline p = TestPipeline.create();
  TupleTag<Long> longTag = new TupleTag<>();
  PCollection<Long> longs = p.apply(GenerateSequence.from(0));
  TupleTag<String> strTag = new TupleTag<>();
  PCollection<String> strs = p.apply(Create.of("foo", "bar"));

  EqualsTester tester = new EqualsTester();
  // Empty tuples in the same pipeline are equal
  tester.addEqualityGroup(PCollectionTuple.empty(p), PCollectionTuple.empty(p));

  tester.addEqualityGroup(
      PCollectionTuple.of(longTag, longs).and(strTag, strs),
      PCollectionTuple.of(longTag, longs).and(strTag, strs));

  tester.addEqualityGroup(PCollectionTuple.of(longTag, longs));
  tester.addEqualityGroup(PCollectionTuple.of(strTag, strs));

  TestPipeline otherPipeline = TestPipeline.create();
  // Empty tuples in different pipelines are not equal
  tester.addEqualityGroup(PCollectionTuple.empty(otherPipeline));
  tester.testEquals();
}
 

@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));
}
 

/** Creates a simple pipeline with a {@link Combine.PerKey} with side inputs. */
private static TestPipeline createCombinePerKeyWithSideInputsPipeline() {
  TestPipeline pipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false);
  PCollection<KV<String, Integer>> input =
      pipeline
          .apply(Create.of(KV.of("key", 1)))
          .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()));
  PCollection<String> sideInput = pipeline.apply(Create.of("side input"));
  PCollectionView<String> sideInputView = sideInput.apply(View.asSingleton());

  input.apply(
      Combine.<String, Integer, Integer>perKey(new SumCombineFnWithContext())
          .withSideInputs(sideInputView));

  return pipeline;
}
 

/** Creates a simple pipeline with a {@link Combine.GroupedValues} with side inputs. */
private static TestPipeline createCombineGroupedValuesWithSideInputsPipeline() {
  TestPipeline pipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false);
  PCollection<KV<String, Integer>> input =
      pipeline
          .apply(Create.of(KV.of("key", 1)))
          .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()));
  PCollection<String> sideInput = pipeline.apply(Create.of("side input"));
  PCollectionView<String> sideInputView = sideInput.apply(View.asSingleton());

  input
      .apply(GroupByKey.create())
      .apply(
          Combine.<String, Integer, Integer>groupedValues(new SumCombineFnWithContext())
              .withSideInputs(sideInputView));

  return pipeline;
}
 

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());
}
 

@org.junit.Test
public void testNest() {
  TestPipeline p = TestPipeline.create().enableAbandonedNodeEnforcement(false);
  PCollection<String> artistText = p.apply("artist", Create.of(artistLinesOfJson)).setCoder(StringUtf8Coder.of());
  PCollection<String> artistCreditNameText = p.apply("artist_credit_name", Create.of(artistCreditLinesOfJson));
  PCollection<String> recordingText = p.apply("recording", Create.of(recordingLinesOfJson)).setCoder(StringUtf8Coder.of());

  PCollection<KV<Long, MusicBrainzDataObject>> artistsById = MusicBrainzTransforms.loadTableFromText(artistText, "artist", "id");

  PCollection<KV<Long, MusicBrainzDataObject>> recordingsByArtistCredit =
      MusicBrainzTransforms.loadTableFromText(recordingText, "recording", "artist_credit");
  PCollection<KV<Long, MusicBrainzDataObject>> artistCreditByArtistCredit =
      MusicBrainzTransforms.loadTableFromText(artistCreditNameText, "artist_credit_name", "artist_credit");

  PCollection<MusicBrainzDataObject> recordingsWithCredits = MusicBrainzTransforms.innerJoin("credited recordings", artistCreditByArtistCredit, recordingsByArtistCredit);
  PCollection<KV<Long, MusicBrainzDataObject>> recordingsJoinedWithCredits =
      MusicBrainzTransforms.by("artist_credit_name_artist", recordingsWithCredits);
  PCollection<MusicBrainzDataObject> artistsWithNestedRecordings = MusicBrainzTransforms.nest(artistsById, recordingsJoinedWithCredits, "recordings");


  PAssert.that(artistsWithNestedRecordings).satisfies( (artistCollection) -> {
    List<MusicBrainzDataObject> theList = new ArrayList<>();
    artistCollection.forEach( (x) -> theList.add(x) );

    assert (theList.size() == 1 );
    assert (((List<MusicBrainzDataObject>) theList.get(0).getColumnValue("artist_recordings")).size() == 448);

    return null;
  });

  p.run();


}
 

@Test
public void testRuntimeValueProviderSubscription() {
  TestPipeline pipeline = TestPipeline.create();
  ValueProvider<String> subscription =
      pipeline.newProvider("projects/project/subscriptions/subscription");
  Read<String> pubsubRead = PubsubIO.readStrings().fromSubscription(subscription);
  pipeline.apply(pubsubRead);
  assertThat(pubsubRead.getSubscriptionProvider(), not(nullValue()));
  assertThat(pubsubRead.getSubscriptionProvider().isAccessible(), is(false));
}
 

@Test
public void testRuntimeValueProviderTopic() {
  TestPipeline pipeline = TestPipeline.create();
  ValueProvider<String> topic = pipeline.newProvider("projects/project/topics/topic");
  Read<String> pubsubRead = PubsubIO.readStrings().fromTopic(topic);
  pipeline.apply(pubsubRead);
  assertThat(pubsubRead.getTopicProvider(), not(nullValue()));
  assertThat(pubsubRead.getTopicProvider().isAccessible(), is(false));
}
 

@Test
public void getInputSucceeds() {
  TestPipeline p = TestPipeline.create();
  PCollection<KV<String, Integer>> input =
      p.apply(
          Create.of(KV.of("foo", 1))
              .withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())));
  PCollection<KV<String, Iterable<Integer>>> grouped = input.apply(GroupByKey.create());
  AppliedPTransform<?, ?, ?> producer = DirectGraphs.getProducer(grouped);
  PTransformReplacement<
          PCollection<KV<String, Integer>>, PCollection<KV<String, Iterable<Integer>>>>
      replacement = factory.getReplacementTransform((AppliedPTransform) producer);
  assertThat(replacement.getInput(), Matchers.<PCollection<?>>equalTo(input));
}