Java TestStream类代码示例

import org.apache.beam.sdk.testing.TestStream; //导入依赖的package包/类
@Test
public void testConvertsBids() throws Exception {
  PCollection<Event> bids = testPipeline.apply(
      TestStream.create(Event.CODER)
          .addElements(new Event(BID))
          .advanceWatermarkToInfinity());

  BeamRecord expectedBidRecord =
      new BeamRecord(
          ModelAdaptersMapping.ADAPTERS.get(Bid.class).getRecordType(),
          ModelAdaptersMapping.ADAPTERS.get(Bid.class).getFieldsValues(BID));

  PAssert
      .that(bids.apply(ToBeamRecord.parDo()))
      .containsInAnyOrder(expectedBidRecord);

  testPipeline.run();
}
 

import org.apache.beam.sdk.testing.TestStream; //导入依赖的package包/类
@Test
public void testConvertsPeople() throws Exception {
  PCollection<Event> people = testPipeline.apply(
      TestStream.create(Event.CODER)
          .addElements(new Event(PERSON))
          .advanceWatermarkToInfinity());

  BeamRecord expectedPersonRecord =
      new BeamRecord(
          ModelAdaptersMapping.ADAPTERS.get(Person.class).getRecordType(),
          ModelAdaptersMapping.ADAPTERS.get(Person.class).getFieldsValues(PERSON));

  PAssert
      .that(people.apply(ToBeamRecord.parDo()))
      .containsInAnyOrder(expectedPersonRecord);

  testPipeline.run();
}
 

import org.apache.beam.sdk.testing.TestStream; //导入依赖的package包/类
@Test
public void testConvertsAuctions() throws Exception {
  PCollection<Event> auctions = testPipeline.apply(
      TestStream.create(Event.CODER)
          .addElements(new Event(AUCTION))
          .advanceWatermarkToInfinity());

  BeamRecord expectedAuctionRecord =
      new BeamRecord(
          ModelAdaptersMapping.ADAPTERS.get(Auction.class).getRecordType(),
          ModelAdaptersMapping.ADAPTERS.get(Auction.class).getFieldsValues(AUCTION));

  PAssert
      .that(auctions.apply(ToBeamRecord.parDo()))
      .containsInAnyOrder(expectedAuctionRecord);

  testPipeline.run();
}
 

import org.apache.beam.sdk.testing.TestStream; //导入依赖的package包/类
@Test
@Category({ValidatesRunner.class, UsesTestStream.class})
public void testReshuffleWithTimestampsStreaming() {
  TestStream<Long> stream =
      TestStream.create(VarLongCoder.of())
          .advanceWatermarkTo(new Instant(0L).plus(Duration.standardDays(48L)))
          .addElements(
              TimestampedValue.of(0L, new Instant(0L)),
              TimestampedValue.of(1L, new Instant(0L).plus(Duration.standardDays(48L))),
              TimestampedValue.of(
                  2L, BoundedWindow.TIMESTAMP_MAX_VALUE.minus(Duration.standardDays(48L))))
          .advanceWatermarkToInfinity();
  PCollection<KV<String, Long>> input =
      pipeline
          .apply(stream).apply(WithKeys.<String, Long>of(""))
          .apply(
              Window.<KV<String, Long>>into(FixedWindows.of(Duration.standardMinutes(10L))));

  PCollection<KV<String, Long>> reshuffled = input.apply(Reshuffle.<String, Long>of());
  PAssert.that(reshuffled.apply(Values.<Long>create())).containsInAnyOrder(0L, 1L, 2L);

  pipeline.run();
}
 

import org.apache.beam.sdk.testing.TestStream; //导入依赖的package包/类
@Test
@Category(NeedsRunner.class)
public void globalWindowNoKeys() {
  PCollection<ValueInSingleWindow<String>> result =
      pipeline
          .apply(
              TestStream.create(StringUtf8Coder.of())
                  .addElements(TimestampedValue.of("dei", new Instant(123L)))
                  .advanceWatermarkToInfinity())
          .apply(Reify.<String>windows());
  PAssert.that(result)
      .containsInAnyOrder(
          ValueInSingleWindow.of(
              "dei", new Instant(123L), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING));
  pipeline.run();
}
 

import org.apache.beam.sdk.testing.TestStream; //导入依赖的package包/类
public RawTestStream(
    RunnerApi.TestStreamPayload payload, RehydratedComponents rehydratedComponents) {
  this.payload = payload;
  this.spec =
      RunnerApi.FunctionSpec.newBuilder()
          .setUrn(TEST_STREAM_TRANSFORM_URN)
          .setPayload(payload.toByteString())
          .build();
  this.rehydratedComponents = rehydratedComponents;

  // Eagerly extract the coder to throw a good exception here
  try {
    this.valueCoder = (Coder<T>) rehydratedComponents.getCoder(payload.getCoderId());
  } catch (IOException exc) {
    throw new IllegalArgumentException(
        String.format(
            "Failure extracting coder with id '%s' for %s",
            payload.getCoderId(), TestStream.class.getSimpleName()),
        exc);
  }
}
 

import org.apache.beam.sdk.testing.TestStream; //导入依赖的package包/类
/**
 * Converts an {@link AppliedPTransform}, which may be a rehydrated transform or an original
 * {@link TestStream}, to a {@link TestStream}.
 */
public static <T> TestStream<T> getTestStream(
    AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>> application)
    throws IOException {
  // For robustness, we don't take this shortcut:
  // if (application.getTransform() instanceof TestStream) {
  //   return application.getTransform()
  // }

  SdkComponents sdkComponents = SdkComponents.create();
  RunnerApi.PTransform transformProto = PTransformTranslation.toProto(application, sdkComponents);
  checkArgument(
      TEST_STREAM_TRANSFORM_URN.equals(transformProto.getSpec().getUrn()),
      "Attempt to get %s from a transform with wrong URN %s",
      TestStream.class.getSimpleName(),
      transformProto.getSpec().getUrn());
  RunnerApi.TestStreamPayload testStreamPayload =
      RunnerApi.TestStreamPayload.parseFrom(transformProto.getSpec().getPayload());

  return (TestStream<T>)
      testStreamFromProtoPayload(
          testStreamPayload, RehydratedComponents.forComponents(sdkComponents.toComponents()));
}
 

import org.apache.beam.sdk.testing.TestStream; //导入依赖的package包/类
@VisibleForTesting
static <T> RunnerApi.TestStreamPayload payloadForTestStream(
    final TestStream<T> testStream, SdkComponents components) throws IOException {
  return payloadForTestStreamLike(
      new TestStreamLike() {
        @Override
        public Coder<T> getValueCoder() {
          return testStream.getValueCoder();
        }

        @Override
        public List<RunnerApi.TestStreamPayload.Event> getEvents() {
          try {
            List<RunnerApi.TestStreamPayload.Event> protoEvents = new ArrayList<>();
            for (TestStream.Event<T> event : testStream.getEvents()) {
              protoEvents.add(eventToProto(event, testStream.getValueCoder()));
            }
            return protoEvents;
          } catch (IOException e) {
            throw new RuntimeException(e);
          }
        }
      },
      components);
}
 

import org.apache.beam.sdk.testing.TestStream; //导入依赖的package包/类
@Test
public void testRegistrarEncodedProto() throws Exception {
  PCollection<String> output = p.apply(testStream);

  AppliedPTransform<PBegin, PCollection<String>, TestStream<String>> appliedTestStream =
      AppliedPTransform.<PBegin, PCollection<String>, TestStream<String>>of(
          "fakeName", PBegin.in(p).expand(), output.expand(), testStream, p);

  SdkComponents components = SdkComponents.create();
  RunnerApi.FunctionSpec spec =
      PTransformTranslation.toProto(appliedTestStream, components).getSpec();

  assertThat(spec.getUrn(), equalTo(TEST_STREAM_TRANSFORM_URN));

  RunnerApi.TestStreamPayload payload = TestStreamPayload.parseFrom(spec.getPayload());

  verifyTestStreamEncoding(
      testStream, payload, RehydratedComponents.forComponents(components.toComponents()));
}
 

import org.apache.beam.sdk.testing.TestStream; //导入依赖的package包/类
private static <T> void verifyTestStreamEncoding(
    TestStream<T> testStream,
    RunnerApi.TestStreamPayload payload,
    RehydratedComponents protoComponents)
    throws Exception {

  // This reverse direction is only valid for Java-based coders
  assertThat(
      protoComponents.getCoder(payload.getCoderId()),
      Matchers.<Coder<?>>equalTo(testStream.getValueCoder()));

  assertThat(payload.getEventsList().size(), equalTo(testStream.getEvents().size()));

  for (int i = 0; i < payload.getEventsList().size(); ++i) {
    assertThat(
        TestStreamTranslation.eventFromProto(payload.getEvents(i), testStream.getValueCoder()),
        equalTo(testStream.getEvents().get(i)));
  }
}
 

import org.apache.beam.sdk.testing.TestStream; //导入依赖的package包/类
/**
 * A test of the {@link CalculateTeamScores} {@link PTransform} when all of the elements arrive
 * on time (ahead of the watermark).
 */
@Test
public void testTeamScoresOnTime() {

  TestStream<GameActionInfo> createEvents = TestStream.create(AvroCoder.of(GameActionInfo.class))
      // Start at the epoch
      .advanceWatermarkTo(baseTime)
      // add some elements ahead of the watermark
      .addElements(event(TestUser.BLUE_ONE, 3, Duration.standardSeconds(3)),
          event(TestUser.BLUE_ONE, 2, Duration.standardMinutes(1)),
          event(TestUser.RED_TWO, 3, Duration.standardSeconds(22)),
          event(TestUser.BLUE_TWO, 5, Duration.standardMinutes(3)))
      // The watermark advances slightly, but not past the end of the window
      .advanceWatermarkTo(baseTime.plus(Duration.standardMinutes(3)))
      // Add some more on time elements
      .addElements(event(TestUser.RED_ONE, 1, Duration.standardMinutes(4)),
          event(TestUser.BLUE_ONE, 2, Duration.standardSeconds(270)))
      // The window should close and emit an ON_TIME pane
      .advanceWatermarkToInfinity();

  PCollection<KV<String, Integer>> teamScores = p.apply(createEvents)
      .apply(new CalculateTeamScores(TEAM_WINDOW_DURATION, ALLOWED_LATENESS));

  String blueTeam = TestUser.BLUE_ONE.getTeam();
  String redTeam = TestUser.RED_ONE.getTeam();
  PAssert.that(teamScores)
      .inOnTimePane(new IntervalWindow(baseTime, TEAM_WINDOW_DURATION))
      .containsInAnyOrder(KV.of(blueTeam, 12), KV.of(redTeam, 4));

  p.run().waitUntilFinish();
}
 

import org.apache.beam.sdk.testing.TestStream; //导入依赖的package包/类
/**
 * A test where elements arrive behind the watermark (late data), but before the end of the
 * window. These elements are emitted on time.
 */
@Test
public void testTeamScoresUnobservablyLate() {

  BoundedWindow window = new IntervalWindow(baseTime, TEAM_WINDOW_DURATION);
  TestStream<GameActionInfo> createEvents = TestStream.create(AvroCoder.of(GameActionInfo.class))
      .advanceWatermarkTo(baseTime)
      .addElements(event(TestUser.BLUE_ONE, 3, Duration.standardSeconds(3)),
          event(TestUser.BLUE_TWO, 5, Duration.standardMinutes(8)),
          event(TestUser.RED_ONE, 4, Duration.standardMinutes(2)),
          event(TestUser.BLUE_ONE, 3, Duration.standardMinutes(5)))
      .advanceWatermarkTo(baseTime.plus(TEAM_WINDOW_DURATION).minus(Duration.standardMinutes(1)))
      // These events are late, but the window hasn't closed yet, so the elements are in the
      // on-time pane
      .addElements(event(TestUser.RED_TWO, 2, Duration.ZERO),
          event(TestUser.RED_TWO, 5, Duration.standardMinutes(1)),
          event(TestUser.BLUE_TWO, 2, Duration.standardSeconds(90)),
          event(TestUser.RED_TWO, 3, Duration.standardMinutes(3)))
      .advanceWatermarkTo(baseTime.plus(TEAM_WINDOW_DURATION).plus(Duration.standardMinutes(1)))
      .advanceWatermarkToInfinity();
  PCollection<KV<String, Integer>> teamScores = p.apply(createEvents)
      .apply(new CalculateTeamScores(TEAM_WINDOW_DURATION, ALLOWED_LATENESS));

  String blueTeam = TestUser.BLUE_ONE.getTeam();
  String redTeam = TestUser.RED_ONE.getTeam();
  // The On Time pane contains the late elements that arrived before the end of the window
  PAssert.that(teamScores)
      .inOnTimePane(window)
      .containsInAnyOrder(KV.of(redTeam, 14), KV.of(blueTeam, 13));

  p.run().waitUntilFinish();
}
 

import org.apache.beam.sdk.testing.TestStream; //导入依赖的package包/类
/**
 * A test where elements arrive beyond the maximum allowed lateness. These elements are dropped
 * within {@link CalculateTeamScores} and do not impact the final result.
 */
@Test
public void testTeamScoresDroppablyLate() {

  BoundedWindow window = new IntervalWindow(baseTime, TEAM_WINDOW_DURATION);
  TestStream<GameActionInfo> infos = TestStream.create(AvroCoder.of(GameActionInfo.class))
      .addElements(event(TestUser.BLUE_ONE, 12, Duration.ZERO),
          event(TestUser.RED_ONE, 3, Duration.ZERO))
      .advanceWatermarkTo(window.maxTimestamp())
      .addElements(event(TestUser.RED_ONE, 4, Duration.standardMinutes(2)),
          event(TestUser.BLUE_TWO, 3, Duration.ZERO),
          event(TestUser.BLUE_ONE, 3, Duration.standardMinutes(3)))
      // Move the watermark to the end of the window to output on time
      .advanceWatermarkTo(baseTime.plus(TEAM_WINDOW_DURATION))
      // Move the watermark past the end of the allowed lateness plus the end of the window
      .advanceWatermarkTo(baseTime.plus(ALLOWED_LATENESS)
          .plus(TEAM_WINDOW_DURATION).plus(Duration.standardMinutes(1)))
      // These elements within the expired window are droppably late, and will not appear in the
      // output
      .addElements(
          event(TestUser.BLUE_TWO, 3, TEAM_WINDOW_DURATION.minus(Duration.standardSeconds(5))),
          event(TestUser.RED_ONE, 7, Duration.standardMinutes(4)))
      .advanceWatermarkToInfinity();
  PCollection<KV<String, Integer>> teamScores = p.apply(infos)
      .apply(new CalculateTeamScores(TEAM_WINDOW_DURATION, ALLOWED_LATENESS));

  String blueTeam = TestUser.BLUE_ONE.getTeam();
  String redTeam = TestUser.RED_ONE.getTeam();
  // Only one on-time pane and no late panes should be emitted
  PAssert.that(teamScores)
      .inWindow(window)
      .containsInAnyOrder(KV.of(redTeam, 7), KV.of(blueTeam, 18));
  // No elements are added before the watermark passes the end of the window plus the allowed
  // lateness, so no refinement should be emitted
  PAssert.that(teamScores).inFinalPane(window).empty();

  p.run().waitUntilFinish();
}
 

import org.apache.beam.sdk.testing.TestStream; //导入依赖的package包/类
@Test
public void testConvertsPriceToEur() throws Exception {
  PCollection<Event> bids = testPipeline.apply(
      TestStream.create(Event.CODER)
          .addElements(new Event(BID1_USD))
          .addElements(new Event(BID2_USD))
          .advanceWatermarkToInfinity());

  PAssert
      .that(bids.apply(new SqlQuery1()))
      .containsInAnyOrder(BID1_EUR_RECORD, BID2_EUR_RECORD);

  testPipeline.run();
}
 

import org.apache.beam.sdk.testing.TestStream; //导入依赖的package包/类
@Test
public void testPassesBidsThrough() throws Exception {
  PCollection<Event> bids = testPipeline.apply(
      TestStream.create(Event.CODER)
          .addElements(new Event(BID1))
          .addElements(new Event(BID2))
          .advanceWatermarkToInfinity());

  PAssert
      .that(bids.apply(new SqlQuery0()))
      .containsInAnyOrder(BID1_RECORD, BID2_RECORD);

  testPipeline.run();
}