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

The following examples show how to use org.apache.beam.sdk.testing.TestPipeline#run() . 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: ExportTimestampTest.java From DataflowTemplates with Apache License 2.0

6 votes

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

Example 2

Source File: ExportTimestampTest.java From DataflowTemplates with Apache License 2.0

6 votes

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

Example 3

Source File: MusicBrainzTransformsTest.java From bigquery-etl-dataflow-sample with Apache License 2.0

5 votes

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


}

Example 4

Source File: MusicBrainzTransformsTest.java From bigquery-etl-dataflow-sample with Apache License 2.0

4 votes

@org.junit.Test
  public void joinArtistCreditsWithRecordings() {

    TestPipeline p = TestPipeline.create().enableAbandonedNodeEnforcement(false);

    PCollection<String> artistCreditText = 
          p.apply("artistCredits", 
                  Create.of(artistCreditLinesOfJson))
                        .setCoder(StringUtf8Coder.of());

    PCollection<KV<Long, MusicBrainzDataObject>> artistCredits =
          MusicBrainzTransforms.loadTableFromText(artistCreditText, "artist_credit_name", "artist_credit");

    PCollection<String> recordingText = 
          p.apply("recordings", 
                  Create.of(recordingLinesOfJson))
                        .setCoder(StringUtf8Coder.of());

    PCollection<KV<Long, MusicBrainzDataObject>> recordings = 
          MusicBrainzTransforms.loadTableFromText(recordingText, "recording", "artist_credit");

    PCollection<MusicBrainzDataObject> joinedRecordings = 
          MusicBrainzTransforms.innerJoin("artist credits with recordings", artistCredits, recordings);

    PCollection<Long> recordingIds = 
        joinedRecordings.apply(MapElements
            .into(new TypeDescriptor<Long>() {})
            .via((MusicBrainzDataObject mbo) -> (Long) mbo.getColumnValue("recording_id"))
        );

    Long bieberRecording = 17069165L;
    Long bieberRecording2 = 15508507L;



    PAssert.that(recordingIds).satisfies((longs) -> {
      List<Long> theList = new ArrayList<Long>();
      longs.forEach(theList::add);
      assert (theList.contains(bieberRecording));
      assert (theList.contains(bieberRecording2));
      return null;
    });


    PCollection<Long> numberJoined = joinedRecordings.apply("count joined recrodings", Count.globally());
    PCollection<Long> numberOfArtistCredits = artistCredits.apply("count artist credits", Count.globally());

    PAssert.thatSingleton(numberJoined).isEqualTo(448L);

    p.run();

//    long joinedRecordingsCount = results.getPCollection(numberJoined).get(0);
//    assert (448 == joinedRecordingsCount);
  }