org.apache.beam.sdk.testing.NeedsRunner Java Examples

@Test
@Category(NeedsRunner.class)
public void loadOneCommitLogFile() {
  PCollection<VersionedEntity> entities =
      pipeline
          .apply(
              "Get CommitLog file patterns",
              Transforms.getCommitLogFilePatterns(commitLogsDir.getAbsolutePath()))
          .apply("Find CommitLogs", Transforms.getFilesByPatterns())
          .apply(Transforms.loadCommitLogsFromFiles());

  InitSqlTestUtils.assertContainsExactlyElementsIn(
      entities,
      KV.of(fakeClock.nowUtc().getMillis() - 2, store.loadAsDatastoreEntity(registry)),
      KV.of(fakeClock.nowUtc().getMillis() - 1, store.loadAsDatastoreEntity(contact)),
      KV.of(fakeClock.nowUtc().getMillis() - 1, store.loadAsDatastoreEntity(domain)));

  pipeline.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testDistinctWithRepresentativeValue() {
  List<KV<String, String>> strings =
      Arrays.asList(KV.of("k1", "v1"), KV.of("k1", "v2"), KV.of("k2", "v1"));

  PCollection<KV<String, String>> input = p.apply(Create.of(strings));

  PCollection<KV<String, String>> output =
      input.apply(
          Distinct.withRepresentativeValueFn(new Keys<String>())
              .withRepresentativeType(TypeDescriptor.of(String.class)));

  PAssert.that(output).satisfies(new Checker());

  p.run();
}
 

/** Basic test of {@link MapElements} with a method reference. */
@Test
@Category(NeedsRunner.class)
public void testMapMethodReference() throws Exception {

  PCollection<Integer> output =
      pipeline
          .apply(Create.of(1, 2, 3))
          .apply(
              MapElements
                  // Note that the type annotation is required.
                  .into(TypeDescriptors.integers())
                  .via(new Doubler()::doubleIt));

  PAssert.that(output).containsInAnyOrder(6, 2, 4);
  pipeline.run();
}
 

@Test
@Category({NeedsRunner.class, UsesSideInputs.class})
public void testSideInputAnnotationFailedValidationSingletonType() {

  final PCollectionView<Integer> sideInput1 =
      pipeline
          .apply("CreateSideInput1", Create.of(2))
          .apply("ViewSideInput1", View.asSingleton());

  // SideInput tag id
  final String sideInputTag1 = "tag1";

  DoFn<Integer, List<Integer>> fn =
      new DoFn<Integer, List<Integer>>() {
        @ProcessElement
        public void processElement(@SideInput(sideInputTag1) String tag1) {}
      };

  thrown.expect(IllegalArgumentException.class);
  PCollection<List<Integer>> output =
      pipeline
          .apply("Create main input", Create.of(2))
          .apply(ParDo.of(fn).withSideInput(sideInputTag1, sideInput1));
  pipeline.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testGlobalAggregation() {
  Collection<Basic> elements =
      ImmutableList.of(
          Basic.of("key1", 1, "value1"),
          Basic.of("key1", 1, "value2"),
          Basic.of("key2", 2, "value3"),
          Basic.of("key2", 2, "value4"));
  PCollection<Long> count =
      pipeline
          .apply(Create.of(elements))
          .apply(Group.<Basic>globally().aggregate(Count.combineFn()));
  PAssert.that(count).containsInAnyOrder(4L);

  pipeline.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testBoundedToUnboundedSourceAdapter() throws Exception {
  long numElements = 100;
  BoundedSource<Long> boundedSource = CountingSource.upTo(numElements);
  UnboundedSource<Long, Checkpoint<Long>> unboundedSource =
      new BoundedToUnboundedSourceAdapter<>(boundedSource);

  PCollection<Long> output = p.apply(Read.from(unboundedSource).withMaxNumRecords(numElements));

  // Count == numElements
  PAssert.thatSingleton(output.apply("Count", Count.globally())).isEqualTo(numElements);
  // Unique count == numElements
  PAssert.thatSingleton(output.apply(Distinct.create()).apply("UniqueCount", Count.globally()))
      .isEqualTo(numElements);
  // Min == 0
  PAssert.thatSingleton(output.apply("Min", Min.globally())).isEqualTo(0L);
  // Max == numElements-1
  PAssert.thatSingleton(output.apply("Max", Max.globally())).isEqualTo(numElements - 1);
  p.run();
}
 

@Test
@Category({NeedsRunner.class, UsesSideInputs.class})
public void testSideInputAnnotationFailedValidationIterableType() {

  final PCollectionView<Iterable<Integer>> sideInput1 =
      pipeline
          .apply("CreateSideInput1", Create.of(2, 1, 0))
          .apply("ViewSideInput1", View.asIterable());

  // SideInput tag id
  final String sideInputTag1 = "tag1";

  DoFn<Integer, List<Integer>> fn =
      new DoFn<Integer, List<Integer>>() {
        @ProcessElement
        public void processElement(@SideInput(sideInputTag1) List<String> tag1) {}
      };

  thrown.expect(IllegalArgumentException.class);
  PCollection<List<Integer>> output =
      pipeline
          .apply("Create main input", Create.of(2))
          .apply(ParDo.of(fn).withSideInput(sideInputTag1, sideInput1));
  pipeline.run();
}
 

/**
 * Tests that {@code AvroIO} can read an upgraded version of an old class, as long as the schema
 * resolution process succeeds. This test covers the case when a new, {@code @Nullable} field
 * has been added.
 *
 * <p>For more information, see http://avro.apache.org/docs/1.7.7/spec.html#Schema+Resolution
 */
@Test
@Category(NeedsRunner.class)
public void testWriteThenReadSchemaUpgrade() throws Throwable {
  List<GenericClass> values =
      ImmutableList.of(new GenericClass(3, "hi"), new GenericClass(5, "bar"));
  File outputFile = tmpFolder.newFile("output.avro");

  writePipeline
      .apply(Create.of(values))
      .apply(
          AvroIO.write(GenericClass.class).to(outputFile.getAbsolutePath()).withoutSharding());
  writePipeline.run();

  List<GenericClassV2> expected =
      ImmutableList.of(new GenericClassV2(3, "hi", null), new GenericClassV2(5, "bar", null));

  PAssert.that(
          readPipeline.apply(
              AvroIO.read(GenericClassV2.class)
                  .withBeamSchemas(withBeamSchemas)
                  .from(outputFile.getAbsolutePath())))
      .containsInAnyOrder(expected);
  readPipeline.run();
}
 

/**
 * Tests that when built with a concrete subclass of {@link SimpleFunction}, the type descriptor
 * of the output reflects its static type.
 */
@Test
@Category(NeedsRunner.class)
public void testSimpleFunctionOutputTypeDescriptor() throws Exception {
  PCollection<String> output =
      pipeline
          .apply(Create.of("hello"))
          .apply(
              MapElements.via(
                  new SimpleFunction<String, String>() {
                    @Override
                    public String apply(String input) {
                      return input;
                    }
                  }));
  assertThat(
      output.getTypeDescriptor(),
      equalTo((TypeDescriptor<String>) new TypeDescriptor<String>() {}));
  assertThat(
      pipeline.getCoderRegistry().getCoder(output.getTypeDescriptor()),
      equalTo(pipeline.getCoderRegistry().getCoder(new TypeDescriptor<String>() {})));

  // Make sure the pipeline runs too
  pipeline.run();
}
 

/**
 * Basic test of {@link MapElements} with a lambda (which is instantiated as a {@link
 * SerializableFunction}).
 */
@Test
@Category(NeedsRunner.class)
public void testMapLambda() throws Exception {

  PCollection<Integer> output =
      pipeline
          .apply(Create.of(1, 2, 3))
          .apply(
              MapElements
                  // Note that the type annotation is required.
                  .into(TypeDescriptors.integers())
                  .via((Integer i) -> i * 2));

  PAssert.that(output).containsInAnyOrder(6, 2, 4);
  pipeline.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testParDoShiftTimestampInvalid() {

  pipeline
      .apply(Create.of(Arrays.asList(3, 42, 6)))
      .apply(ParDo.of(new TestOutputTimestampDoFn<>()))
      .apply(
          ParDo.of(
              new TestShiftTimestampDoFn<>(
                  Duration.millis(1000), // allowed skew = 1 second
                  Duration.millis(-1001))))
      .apply(ParDo.of(new TestFormatTimestampDoFn<>()));

  thrown.expect(RuntimeException.class);
  thrown.expectMessage("Cannot output with timestamp");
  thrown.expectMessage(
      "Output timestamps must be no earlier than the timestamp of the current input");
  thrown.expectMessage("minus the allowed skew (1 second).");
  pipeline.run();
}
 

/** Basic test of {@link MapElements} with a {@link SimpleFunction}. */
@Test
@Category(NeedsRunner.class)
public void testMapSimpleFunction() throws Exception {
  PCollection<Integer> output =
      pipeline
          .apply(Create.of(1, 2, 3))
          .apply(
              MapElements.via(
                  new SimpleFunction<Integer, Integer>() {
                    @Override
                    public Integer apply(Integer input) {
                      return -input;
                    }
                  }));

  PAssert.that(output).containsInAnyOrder(-2, -1, -3);
  pipeline.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testUnionAll() {

  PAssert.that(first.apply("strings", Sets.unionAll(second)))
      .containsInAnyOrder(
          "a", "a", "a", "a", "a", "b", "b", "b", "b", "b", "c", "c", "d", "d", "d", "d", "e",
          "e", "f", "f", "g", "g", "h", "h");

  PCollection<Row> results = firstRows.apply("rows", Sets.unionAll(secondRows));

  PAssert.that(results)
      .containsInAnyOrder(
          toRows(
              "a", "a", "a", "a", "a", "b", "b", "b", "b", "b", "c", "c", "d", "d", "d", "d", "e",
              "e", "f", "f", "g", "g", "h", "h"));

  assertEquals(schema, results.getSchema());

  p.run();
}
 

/**
 * Test {@link TransformTextViaJavascript} passes through data when null ValueProvider as args.
 * when hasInvocable returns false.
 */
@Test
@Category(NeedsRunner.class)
public void testDoFnPassthroughNullValueProvider() {
  List<String> inJson = Arrays.asList("{\"answerToLife\":    42}");

  PCollection<String> transformedJson =
      pipeline
          .apply("Create", Create.of(inJson))
          .apply(
              TransformTextViaJavascript.newBuilder()
                  .setFunctionName(null)
                  .setFileSystemPath(null)
                  .build());

  PAssert.that(transformedJson).containsInAnyOrder(inJson);

  pipeline.run();
}
 

/**
 * Read a ZIP compressed file containing data, multiple empty entries, and then more data. We
 * expect just the data back.
 */
@Test
@Category(NeedsRunner.class)
public void testZipCompressedReadWithComplexEmptyAndPresentEntries() throws Exception {
  File file =
      createZipFile(
          new ArrayList<>(),
          tempFolder,
          "complex empty and present entries",
          new String[] {"cat"},
          new String[] {},
          new String[] {},
          new String[] {"dog"});

  assertReadingCompressedFileMatchesExpected(file, ZIP, Arrays.asList("cat", "dog"), p);
  p.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testAggregateByMultipleFields() {
  Collection<Aggregate> elements =
      ImmutableList.of(
          Aggregate.of(1, 1, 2),
          Aggregate.of(2, 1, 3),
          Aggregate.of(3, 2, 4),
          Aggregate.of(4, 2, 5));

  List<String> fieldNames = Lists.newArrayList("field1", "field2");
  PCollection<Row> aggregate =
      pipeline
          .apply(Create.of(elements))
          .apply(
              Group.<Aggregate>globally()
                  .aggregateFields(fieldNames, new MultipleFieldCombineFn(), "field1+field2"));

  Schema outputSchema = Schema.builder().addInt64Field("field1+field2").build();
  Row expectedRow = Row.withSchema(outputSchema).addValues(16L).build();
  PAssert.that(aggregate).containsInAnyOrder(expectedRow);

  pipeline.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testDroppedPartition() {

  // Compute the set of integers either 1 or 2 mod 3, the hard way.
  PCollectionList<Integer> outputs =
      pipeline
          .apply(Create.of(2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12))
          .apply(Partition.of(3, new ModFn()));

  List<PCollection<Integer>> outputsList = new ArrayList<>(outputs.getAll());
  outputsList.remove(0);
  outputs = PCollectionList.of(outputsList);
  assertTrue(outputs.size() == 2);

  PCollection<Integer> output = outputs.apply(Flatten.pCollections());
  PAssert.that(output).containsInAnyOrder(2, 4, 5, 7, 8, 10, 11);
  pipeline.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testPartitionWithSideInputs() {

  PCollectionView<Integer> gradesView =
      pipeline.apply("grades", Create.of(50)).apply(View.asSingleton());

  Create.Values<Integer> studentsPercentage = Create.of(5, 45, 90, 29, 55, 65);
  PCollectionList<Integer> studentsGrades =
      pipeline
          .apply(studentsPercentage)
          .apply(
              Partition.of(
                  2,
                  ((elem, numPartitions, ct) -> {
                    Integer grades = ct.sideInput(gradesView);
                    return elem < grades ? 0 : 1;
                  }),
                  Requirements.requiresSideInputs(gradesView)));

  assertTrue(studentsGrades.size() == 2);
  PAssert.that(studentsGrades.get(0)).containsInAnyOrder(5, 29, 45);
  PAssert.that(studentsGrades.get(1)).containsInAnyOrder(55, 65, 90);
  pipeline.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testComposeBasicSerializableFunction() throws Exception {
  PCollection<Integer> output =
      pipeline
          .apply(Create.of(1, 2, 3))
          .apply(
              PTransform.compose(
                  (PCollection<Integer> numbers) -> {
                    PCollection<Integer> inverted =
                        numbers.apply(MapElements.into(integers()).via(input -> -input));
                    return PCollectionList.of(numbers)
                        .and(inverted)
                        .apply(Flatten.pCollections());
                  }));

  PAssert.that(output).containsInAnyOrder(-2, -1, -3, 2, 1, 3);
  pipeline.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testUnionCollectionList() {
  PCollection<String> third = p.apply("third", Create.of(Arrays.asList("a", "k", "k")));
  PCollection<Row> thirdRows = p.apply("thirdRows", Create.of(toRows("a", "k", "k")));

  PAssert.that(
          PCollectionList.of(first)
              .and(second)
              .and(third)
              .apply("stringsCols", Sets.unionDistinct()))
      .containsInAnyOrder("a", "b", "c", "d", "e", "f", "g", "h", "k");

  PCollection<Row> results =
      PCollectionList.of(firstRows)
          .and(secondRows)
          .and(thirdRows)
          .apply("rowCols", Sets.unionDistinct());

  PAssert.that(results).containsInAnyOrder(toRows("a", "b", "c", "d", "e", "f", "g", "h", "k"));

  assertEquals(schema, results.getSchema());

  p.run();
}
 

@Test
@Category(NeedsRunner.class)
public void getCommitLogFilePatterns() {
  PCollection<String> patterns =
      pipeline.apply(
          "Get CommitLog file patterns",
          Transforms.getCommitLogFilePatterns(commitLogsDir.getAbsolutePath()));

  ImmutableList<String> expectedPatterns =
      ImmutableList.of(commitLogsDir.getAbsolutePath() + "/commit_diff_until_*");

  PAssert.that(patterns).containsInAnyOrder(expectedPatterns);

  pipeline.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testModPartitionWithLambda() {

  PCollectionList<Integer> outputs =
      pipeline
          .apply(Create.of(1, 2, 4, 5))
          .apply(Partition.of(3, (element, numPartitions) -> element % numPartitions));
  assertEquals(3, outputs.size());
  PAssert.that(outputs.get(0)).empty();
  PAssert.that(outputs.get(1)).containsInAnyOrder(1, 4);
  PAssert.that(outputs.get(2)).containsInAnyOrder(2, 5);
  pipeline.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testNoFilterByPredicateWithLambda() {

  PCollection<Integer> output = p.apply(Create.of(1, 2, 4, 5)).apply(Filter.by(i -> false));

  PAssert.that(output).empty();
  p.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testQuantilesPerKey() {
  PCollection<KV<String, Integer>> input = createInputTable(p);
  PCollection<KV<String, List<Integer>>> quantiles =
      input.apply(ApproximateQuantiles.perKey(2));

  PAssert.that(quantiles)
      .containsInAnyOrder(KV.of("a", Arrays.asList(1, 3)), KV.of("b", Arrays.asList(1, 100)));
  p.run();
}
 

/** Test successful conversion of JSON messages with a user provided host. */
@Test
@Category(NeedsRunner.class)
public void testFailsafeStringToSplunkEventValidHost() {

  FailsafeElement<String, String> input =
      FailsafeElement.of(
          "",
          "{\n"
              + "\t\"name\": \"Jim\",\n"
              + "\t\"_metadata\": {\"host\": \"test-host\"}\n"
              + "}");

  pipeline.getCoderRegistry().registerCoderForClass(SplunkEvent.class, SplunkEventCoder.of());

  PCollectionTuple tuple =
      pipeline
          .apply(
              Create.of(input)
                  .withCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))
          .apply(
              SplunkConverters.failsafeStringToSplunkEvent(
                  SPLUNK_EVENT_OUT, SPLUNK_EVENT_DEADLETTER_OUT));

  PAssert.that(tuple.get(SPLUNK_EVENT_DEADLETTER_OUT)).empty();
  PAssert.that(tuple.get(SPLUNK_EVENT_OUT))
      .containsInAnyOrder(
          SplunkEvent.newBuilder().withEvent("{\"name\":\"Jim\"}").withHost("test-host").build());

  pipeline.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testUnboundedSourceWithRate() {

  Duration period = Duration.millis(5);
  long numElements = 1000L;

  PCollection<Long> input =
      p.apply(
          Read.from(
                  CountingSource.createUnboundedFrom(0)
                      .withTimestampFn(new ValueAsTimestampFn())
                      .withRate(1, period))
              .withMaxNumRecords(numElements));
  addCountingAsserts(input, numElements);

  PCollection<Long> diffs =
      input
          .apply("TimestampDiff", ParDo.of(new ElementValueDiff()))
          .apply("DistinctTimestamps", Distinct.create());
  // This assert also confirms that diffs only has one unique value.
  PAssert.thatSingleton(diffs).isEqualTo(0L);

  Instant started = Instant.now();
  p.run();
  Instant finished = Instant.now();
  Duration expectedDuration = period.multipliedBy((int) numElements);
  assertThat(started.plus(expectedDuration).isBefore(finished), is(true));
}
 

@Test
@Category(NeedsRunner.class)
public void testFilterByProcessFunction() {
  PCollection<Integer> output =
      p.apply(Create.of(1, 2, 3, 4, 5, 6, 7)).apply(Filter.by(new EvenProcessFn()));

  PAssert.that(output).containsInAnyOrder(2, 4, 6);
  p.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testToStringOf() {
  Integer[] ints = {1, 2, 3, 4, 5};
  String[] strings = {"1", "2", "3", "4", "5"};
  PCollection<Integer> input = p.apply(Create.of(Arrays.asList(ints)));
  PCollection<String> output = input.apply(ToString.elements());
  PAssert.that(output).containsInAnyOrder(strings);
  p.run();
}
 

@Test
@Category(NeedsRunner.class)
public void recursivelyAddNestedArrayField() {
  Schema schema = Schema.builder().build();
  Row row = Row.withSchema(schema).build();
  PCollection<Row> added =
      pipeline
          .apply(Create.of(row).withRowSchema(schema))
          .apply(
              AddFields.<Row>create()
                  .field("array[].field1", FieldType.STRING)
                  .field("array[].field2", Schema.FieldType.INT32)
                  .field("array[].field3", Schema.FieldType.array(Schema.FieldType.STRING)));

  Schema expectedNestedSchema =
      Schema.builder()
          .addNullableField("field1", FieldType.STRING)
          .addNullableField("field2", Schema.FieldType.INT32)
          .addNullableField("field3", Schema.FieldType.array(Schema.FieldType.STRING))
          .build();
  Schema expectedSchema =
      Schema.builder()
          .addNullableField(
              "array",
              Schema.FieldType.array(
                  Schema.FieldType.row(expectedNestedSchema).withNullable(true)))
          .build();
  assertEquals(expectedSchema, added.getSchema());

  Row expected = Row.withSchema(expectedSchema).addValue(Collections.emptyList()).build();
  PAssert.that(added).containsInAnyOrder(expected);
  pipeline.run();
}
 

@Test
@Category(NeedsRunner.class)
public void testMergePartialGloballyForMergeWithEmptySketch() {
  PCollection<byte[]> result =
      p.apply(Create.of(LONGS_SKETCH, EMPTY_SKETCH)).apply(HllCount.MergePartial.globally());

  PAssert.thatSingleton(result).isEqualTo(LONGS_SKETCH);
  p.run();
}