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

The following examples show how to use org.apache.beam.sdk.testing.TestStream#Event . 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: TestStreamTranslation.java    From beam with Apache License 2.0 6 votes vote down vote up 
/** 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();
}
Example 2
Source File: TestStreamTranslation.java    From beam with Apache License 2.0 5 votes vote down vote up 
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);
}
Example 3
Source File: TestStreamTranslation.java    From beam with Apache License 2.0 5 votes vote down vote up 
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()));
  }
}
Example 4
Source File: TestStreamTranslation.java    From beam with Apache License 2.0 5 votes vote down vote up 
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()));
  }
}
Example 5
Source File: TestStreamSource.java    From beam with Apache License 2.0 5 votes vote down vote up 
@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);
      }
    }
  }
}
Example 6
Source File: DataflowPipelineTranslator.java    From beam with Apache License 2.0 4 votes vote down vote up 
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);
  }
}