java - 如何对基于 KinesisRecord 的 DoFn 进行单元测试？

Question

我正在开始一个从 AWS Kinesis 读取数据的 Beam 项目，因此我有一个简单的 DoFn 来接受 KinesisRecord 并记录内容。我想编写一个单元测试来运行这个 DoFn 并证明它有效。不过，事实证明，使用 KinesisRecord 进行单元测试具有挑战性。

当我尝试仅使用 Create.of(testKinesisRecord) 时收到此错误:

java.lang.IllegalArgumentException: Unable to infer a coder and no Coder was specified. Please set a coder by invoking Create.withCoder() explicitly  or a schema by invoking Create.withSchema().

我尝试按照错误提示显式使用“withCoder”提供 KinesisRecordCoder，但它是一个私有(private)类。也许还有另一种方法来对 DoFn 进行单元测试？

测试代码:

public class MyProjectTests {
    @Rule
    public TestPipeline p = TestPipeline.create();

    @Test
    public void testPoC() {
        var testKinesisRecord = new KinesisRecord(
                ByteBuffer.wrap("SomeData".getBytes()),
                "seq01",
                12,
                "pKey",
                Instant.now().minus(Duration.standardHours(4)),
                Instant.now(),
                "MyStream",
                "shard-001"
        );


        PCollection<Void> output =
                p.apply(Create.of(testKinesisRecord))
                        .apply(ParDo.of(new MyProject.PrintRecordFn()));

        var result = p.run();
        result.waitUntilFinish();
        result.metrics().allMetrics().getCounters().forEach(longMetricResult -> {
            Assertions.assertEquals(1, longMetricResult.getCommitted().intValue());
        });
    }
}

DoFn 代码:

  static class PrintRecordFn extends DoFn<KinesisRecord, Void> {
    private static final Logger LOG = LoggerFactory.getLogger(PrintRecordFn.class);
    private final Counter items = Metrics.counter(PrintRecordFn.class, "itemsProcessed");

    @ProcessElement
    public void processElement(@Element KinesisRecord element) {
      items.inc();

      LOG.info("Stream: `{}` Shard: `{}` Arrived at `{}`\nData: {}",
              element.getStreamName(),
              element.getShardId(),
              element.getApproximateArrivalTimestamp(),
              element.getDataAsBytes());
    }
  }

Answer 1

KinesisRecordCoder 应该用于内部目的，因此它被设为包私有(private)。同时，您可以提供自定义的AWSClientsProvider并使用它来生成测试数据。作为示例，请查看 KinesisMockReadTest和定制Provider