Java Code Examples for org.apache.beam.sdk.testing.TestStream#Event

/** Produces a {@link RunnerApi.TestStreamPayload} from a {@link TestStream}. */
static <T> RunnerApi.TestStreamPayload payloadForTestStream(
    final TestStream<T> transform, SdkComponents components) throws IOException {
  List<RunnerApi.TestStreamPayload.Event> protoEvents = new ArrayList<>();
  try {
    for (TestStream.Event<T> event : transform.getEvents()) {
      protoEvents.add(eventToProto(event, transform.getValueCoder()));
    }
  } catch (IOException e) {
    throw new RuntimeException(e);
  }

  return RunnerApi.TestStreamPayload.newBuilder()
      .setCoderId(components.registerCoder(transform.getValueCoder()))
      .addAllEvents(protoEvents)
      .build();
}
 

public static TestStream<?> testStreamFromProtoPayload(
    RunnerApi.TestStreamPayload testStreamPayload, RehydratedComponents components)
    throws IOException {

  Coder<Object> coder = (Coder<Object>) components.getCoder(testStreamPayload.getCoderId());

  List<TestStream.Event<Object>> events = new ArrayList<>();

  for (RunnerApi.TestStreamPayload.Event event : testStreamPayload.getEventsList()) {
    events.add(eventFromProto(event, coder));
  }
  return TestStream.fromRawEvents(coder, events);
}
 

static <T> RunnerApi.TestStreamPayload.Event eventToProto(
    TestStream.Event<T> event, Coder<T> coder) throws IOException {
  switch (event.getType()) {
    case WATERMARK:
      return RunnerApi.TestStreamPayload.Event.newBuilder()
          .setWatermarkEvent(
              RunnerApi.TestStreamPayload.Event.AdvanceWatermark.newBuilder()
                  .setNewWatermark(
                      ((TestStream.WatermarkEvent<T>) event).getWatermark().getMillis()))
          .build();

    case PROCESSING_TIME:
      return RunnerApi.TestStreamPayload.Event.newBuilder()
          .setProcessingTimeEvent(
              RunnerApi.TestStreamPayload.Event.AdvanceProcessingTime.newBuilder()
                  .setAdvanceDuration(
                      ((TestStream.ProcessingTimeEvent<T>) event)
                          .getProcessingTimeAdvance()
                          .getMillis()))
          .build();

    case ELEMENT:
      RunnerApi.TestStreamPayload.Event.AddElements.Builder builder =
          RunnerApi.TestStreamPayload.Event.AddElements.newBuilder();
      for (TimestampedValue<T> element : ((TestStream.ElementEvent<T>) event).getElements()) {
        builder.addElements(
            RunnerApi.TestStreamPayload.TimestampedElement.newBuilder()
                .setTimestamp(element.getTimestamp().getMillis())
                .setEncodedElement(
                    ByteString.copyFrom(
                        CoderUtils.encodeToByteArray(coder, element.getValue()))));
      }
      return RunnerApi.TestStreamPayload.Event.newBuilder().setElementEvent(builder).build();
    default:
      throw new IllegalArgumentException(
          String.format(
              "Unsupported type of %s: %s",
              TestStream.Event.class.getCanonicalName(), event.getType()));
  }
}
 

static <T> TestStream.Event<T> eventFromProto(
    RunnerApi.TestStreamPayload.Event protoEvent, Coder<T> coder) throws IOException {
  switch (protoEvent.getEventCase()) {
    case WATERMARK_EVENT:
      return TestStream.WatermarkEvent.advanceTo(
          new Instant(protoEvent.getWatermarkEvent().getNewWatermark()));
    case PROCESSING_TIME_EVENT:
      return TestStream.ProcessingTimeEvent.advanceBy(
          Duration.millis(protoEvent.getProcessingTimeEvent().getAdvanceDuration()));
    case ELEMENT_EVENT:
      List<TimestampedValue<T>> decodedElements = new ArrayList<>();
      for (RunnerApi.TestStreamPayload.TimestampedElement element :
          protoEvent.getElementEvent().getElementsList()) {
        decodedElements.add(
            TimestampedValue.of(
                CoderUtils.decodeFromByteArray(coder, element.getEncodedElement().toByteArray()),
                new Instant(element.getTimestamp())));
      }
      return TestStream.ElementEvent.add(decodedElements);
    case EVENT_NOT_SET:
    default:
      throw new IllegalArgumentException(
          String.format(
              "Unsupported type of %s: %s",
              RunnerApi.TestStreamPayload.Event.class.getCanonicalName(),
              protoEvent.getEventCase()));
  }
}
 

@Override
public void run(SourceContext<WindowedValue<T>> ctx) throws CoderException {
  TestStream<T> testStream = testStreamDecoder.apply(payload);
  List<TestStream.Event<T>> events = testStream.getEvents();

  for (int eventId = 0; isRunning && eventId < events.size(); eventId++) {
    TestStream.Event<T> event = events.get(eventId);

    synchronized (ctx.getCheckpointLock()) {
      if (event instanceof TestStream.ElementEvent) {
        for (TimestampedValue<T> element : ((TestStream.ElementEvent<T>) event).getElements()) {
          Instant timestamp = element.getTimestamp();
          WindowedValue<T> value =
              WindowedValue.of(
                  element.getValue(), timestamp, GlobalWindow.INSTANCE, PaneInfo.NO_FIRING);
          ctx.collectWithTimestamp(value, timestamp.getMillis());
        }
      } else if (event instanceof TestStream.WatermarkEvent) {
        long millis = ((TestStream.WatermarkEvent<T>) event).getWatermark().getMillis();
        ctx.emitWatermark(new Watermark(millis));
      } else if (event instanceof TestStream.ProcessingTimeEvent) {
        // There seems to be no clean way to implement this
        throw new UnsupportedOperationException(
            "Advancing Processing time is not supported by the Flink Runner.");
      } else {
        throw new IllegalStateException("Unknown event type " + event);
      }
    }
  }
}
 

private <T> void translateTyped(TestStream<T> transform, TranslationContext context) {
  try {
    StepTranslationContext stepContext = context.addStep(transform, "ParallelRead");
    String ptransformId =
        context.getSdkComponents().getPTransformIdOrThrow(context.getCurrentTransform());
    stepContext.addInput(PropertyNames.SERIALIZED_FN, ptransformId);
    stepContext.addInput(PropertyNames.FORMAT, "test_stream");
    RunnerApi.TestStreamPayload.Builder payloadBuilder =
        RunnerApi.TestStreamPayload.newBuilder();
    for (TestStream.Event event : transform.getEvents()) {
      if (event instanceof TestStream.ElementEvent) {
        RunnerApi.TestStreamPayload.Event.AddElements.Builder addElementsBuilder =
            RunnerApi.TestStreamPayload.Event.AddElements.newBuilder();
        Iterable<TimestampedValue<T>> elements =
            ((TestStream.ElementEvent) event).getElements();
        for (TimestampedValue<T> element : elements) {
          addElementsBuilder.addElements(
              RunnerApi.TestStreamPayload.TimestampedElement.newBuilder()
                  .setEncodedElement(
                      ByteString.copyFrom(
                          CoderUtils.encodeToByteArray(
                              transform.getValueCoder(), element.getValue())))
                  .setTimestamp(element.getTimestamp().getMillis() * 1000));
        }
        payloadBuilder.addEventsBuilder().setElementEvent(addElementsBuilder);
      } else if (event instanceof TestStream.WatermarkEvent) {
        payloadBuilder
            .addEventsBuilder()
            .setWatermarkEvent(
                RunnerApi.TestStreamPayload.Event.AdvanceWatermark.newBuilder()
                    .setNewWatermark(
                        ((TestStream.WatermarkEvent) event).getWatermark().getMillis()
                            * 1000));
      } else if (event instanceof TestStream.ProcessingTimeEvent) {
        payloadBuilder
            .addEventsBuilder()
            .setProcessingTimeEvent(
                RunnerApi.TestStreamPayload.Event.AdvanceProcessingTime.newBuilder()
                    .setAdvanceDuration(
                        ((TestStream.ProcessingTimeEvent) event)
                                .getProcessingTimeAdvance()
                                .getMillis()
                            * 1000));
      }
    }
    stepContext.addInput(
        PropertyNames.SERIALIZED_TEST_STREAM,
        byteArrayToJsonString(payloadBuilder.build().toByteArray()));
    stepContext.addOutput(PropertyNames.OUTPUT, context.getOutput(transform));
  } catch (Exception e) {
    throw new RuntimeException(e);
  }
}