Java Code Examples for org.apache.beam.sdk.transforms.DoFnTester#of()

@Test
public void testHierarchical_TFD2119_B7_HasAtLeastOneSubrecordWithSubSubRecordValueGt10() throws Exception {
    DoFnTester<IndexedRecord, IndexedRecord> fnTester = DoFnTester.of( //
            new FilterRowDoFn(addCriteria(null, //
                    ".b1{.b2.value > 10}", //
                    ConditionsRowConstant.Function.COUNT, //
                    ConditionsRowConstant.Operator.GREATER, //
                    "0") //
            ));

    List<IndexedRecord> output = fnTester.processBundle(inputB);

    for (IndexedRecord main : output) {
        boolean atLeastOne = false;
        for (IndexedRecord subrecord : getSubrecords(main)) {
            for (IndexedRecord subsubrecord : getSubrecords(subrecord)) {
                if ((double) subsubrecord.get(2) > 10)
                    atLeastOne = true;
            }
        }
        assertThat(main.toString(), atLeastOne, is(true));
    }
    assertThat(output, hasSize(311));
}
 

@Test
public void selectSimpleElementsWithEmptyValues() throws Exception {
    FieldSelectorProperties properties = addSelector(null, "aOutput", "a");
    properties = addSelector(properties, "cOutput", "c");
    properties = addSelector(properties, "aSecondOutput", "a");
    properties = addSelector(properties, "bOutput", "b");

    FieldSelectorDoFn function = new FieldSelectorDoFn().withProperties(properties);
    DoFnTester<IndexedRecord, IndexedRecord> fnTester = DoFnTester.of(function);
    List<IndexedRecord> outputs = fnTester.processBundle(inputRecordWithEmptyValue);

    assertEquals(1, outputs.size());
    List<Field> fields = outputs.get(0).getSchema().getFields();
    assertEquals(4, fields.size());
    assertEquals("aOutput", fields.get(0).name());
    assertEquals("aaa", outputs.get(0).get(0));
    assertEquals("cOutput", fields.get(1).name());
    assertNull(outputs.get(0).get(1));
    assertEquals("aSecondOutput", fields.get(2).name());
    assertEquals("aaa", outputs.get(0).get(2));
    assertEquals("bOutput", fields.get(3).name());
    assertEquals("BBB", outputs.get(0).get(3));
}
 

@Test
public void testGetAlleleFreq() throws Exception {
  DoFnTester<Variant, KV<Position, AlleleFreq>> getAlleleFreq = DoFnTester.of(
      new GetAlleleFreq());
  Position pos = Position.newBuilder()
      .setReferenceName("1")
      .setPosition(123L)
      .build();
  Variant.Builder vBuild = Variant.newBuilder()
      .setReferenceName("1")
      .setStart(123L)
      .setReferenceBases("C")
      .addAlternateBases("T");
  vBuild.getMutableInfo().put("AF", ListValue.newBuilder()
      .addValues(Value.newBuilder().setStringValue("0.25").build()).build());
  AlleleFreq af = new AlleleFreq();
  af.setAltBases(Lists.newArrayList("T"));
  af.setRefBases("C");
  af.setRefFreq(0.25);
  Assert.assertThat(getAlleleFreq.processBundle(vBuild.build()),
      CoreMatchers.hasItems(KV.of(pos, af)));
}
 

/**
 * Input parent record: {@link NormalizeDoFnTest#inputParentRecord}
 *
 * Normalize simple field: `a`
 *
 * Expected normalized results of the field `a`:
 *
 * {"a": "aaa", "b": {"x": "x1;x2", "y": {"d": {"j": [{"l": "l1"}, {"l": "l2"}], "k": "k1;k2"}, "e": "e"}}, "c":
 * {"f": "f", "g": [{"h": "h1", "i": "i2"}, {"h": "h2", "i": "i1"}]}, "m": ["m1", "m2", "m3"]}
 *
 * @throws Exception
 */
@Test
public void testNormalizeSimpleFields_a() throws Exception {
    NormalizeProperties properties = new NormalizeProperties("test");
    properties.init();
    properties.schemaListener.afterSchema();
    properties.isList.setValue(false);
    properties.trim.setValue(true);
    properties.discardTrailingEmptyStr.setValue(true);

    // Normalize `a` simple field
    properties.columnToNormalize.setValue("a");

    NormalizeDoFn function = new NormalizeDoFn().withProperties(properties);
    DoFnTester<IndexedRecord, IndexedRecord> fnTester = DoFnTester.of(function);
    List<IndexedRecord> outputs = fnTester.processBundle(inputParentRecord);
    Assert.assertEquals(1, outputs.size());
    GenericRecord outputRecord = (GenericRecord) outputs.get(0);
    Assert.assertEquals(inputParentRecord.toString(), outputRecord.toString());
    Assert.assertEquals(inputParentRecord.getSchema().toString(), outputRecord.getSchema().toString());
}
 

@Test
public void test_Map_ApplyATransformation() throws Exception {

    PythonRowProperties properties = new PythonRowProperties("test");
    properties.init();
    properties.mapType.setValue(MapType.MAP);

    StringBuilder sb = new StringBuilder();
    sb.append("output = input\n");
    sb.append("output['a1'] = \"rootdata2\"\n");
    sb.append("output['B']['b1'] = \"subdatabefore\"\n");
    sb.append("output['B']['C']['c1'] = \"subsubdatabefore\"\n");
    sb.append("output['B']['C']['c2'] = 33\n");
    sb.append("output['B']['C']['c3'] = 55l\n");
    sb.append("output['B']['b2'] = \"subdataend\"\n");
    properties.pythonCode.setValue(sb.toString());
    PythonRowDoFn function = new PythonRowDoFn();
    assertEquals(ValidationResult.OK, function.initialize(null, properties));
    DoFnTester<IndexedRecord, IndexedRecord> fnTester = DoFnTester.of(function);
    List<IndexedRecord> outputs = fnTester.processBundle(inputIndexedRecord);
    assertEquals(1, outputs.size());

    GenericRecord outputRecord = (GenericRecord) outputs.get(0);
    compareRecords(outputIndexedRecord, outputRecord);
}
 

/**
 * Input parent record: {@link NormalizeDoFnTest#inputParentRecord}
 *
 * Normalize complex field: `b`
 *
 * Expected: no change
 *
 * @throws Exception
 */
@Test
public void testNormalizeComplexFields_b() throws Exception {
    NormalizeProperties properties = new NormalizeProperties("test");
    properties.init();
    properties.schemaListener.afterSchema();

    // Normalize `b` complex field
    properties.columnToNormalize.setValue("b");

    NormalizeDoFn function = new NormalizeDoFn().withProperties(properties);
    DoFnTester<IndexedRecord, IndexedRecord> fnTester = DoFnTester.of(function);
    List<IndexedRecord> outputs = fnTester.processBundle(inputParentRecord);
    Assert.assertEquals(1, outputs.size());
    GenericRecord outputRecord = (GenericRecord) outputs.get(0);
    Assert.assertEquals(inputParentRecord.toString(), outputRecord.toString());
    Assert.assertEquals(inputParentRecord.getSchema().toString(), outputRecord.getSchema().toString());
}
 

@Test
public void testHierarchical_TFD2119_B6_AllSubRecordsWithId1Or2HasValueGt10() throws Exception {
    DoFnTester<IndexedRecord, IndexedRecord> fnTester = DoFnTester.of( //
            new FilterRowDoFn(addCriteria(null, //
                    ".b1{.id == 1 || .id == 2}.value", //
                    null, //
                    ConditionsRowConstant.Operator.GREATER, //
                    "10") //
            ));

    List<IndexedRecord> output = fnTester.processBundle(inputB);

    for (IndexedRecord main : output) {
        boolean atLeastOne = false;
        for (IndexedRecord subrecord : getSubrecords(main)) {
            int id = (int) subrecord.get(0);
            if (id == 1 || id == 2) {
                atLeastOne = true;
                assertThat(main.toString(), (double) subrecord.get(2), greaterThan(10d));
            }
        }
        assertThat(main.toString(), atLeastOne, is(true));
    }
    assertThat(output, hasSize(42));
}
 

/**
 * Test the {@link LimitDoFn}.
 */
@Test
public void testDoFn() throws Exception {

    // Process the same number of records than the limit
    DoFnTester<IndexedRecord, IndexedRecord> fnTester1 = DoFnTester.of(createLimitFunction(3L));
    List<IndexedRecord> outputs1 = fnTester1.processBundle(inputSimpleRecord, inputSimpleRecord, inputSimpleRecord);
    assertEquals(3, outputs1.size());

    // Process more records than the limit
    DoFnTester<IndexedRecord, IndexedRecord> fnTester2 = DoFnTester.of(createLimitFunction(2L));
    List<IndexedRecord> outputs2 = fnTester2.processBundle(inputSimpleRecord, inputSimpleRecord, inputSimpleRecord);
    assertEquals(2, outputs2.size());

    // Process less records than the limit
    DoFnTester<IndexedRecord, IndexedRecord> fnTester3 = DoFnTester.of(createLimitFunction(4L));
    List<IndexedRecord> outputs3 = fnTester3.processBundle(inputSimpleRecord, inputSimpleRecord, inputSimpleRecord);
    assertEquals(3, outputs3.size());
}
 

@Test
public void testHierarchical_TFD2119_A1_TopLevel() throws Exception {
    DoFnTester<IndexedRecord, IndexedRecord> fnTester = DoFnTester.of( //
            new FilterRowDoFn(addCriteria(null, //
                    ".id", //
                    null, //
                    ConditionsRowConstant.Operator.EQUAL, //
                    "1") //
            ));

    List<IndexedRecord> output = fnTester.processBundle(inputA);
    for (IndexedRecord main : output) {
        assertThat(main.toString(), main.get(0), is((Object) 1));
    }
    assertThat(output, hasSize(103));
}
 

@Test
public void test_FilterBetween() throws Exception {
    FilterRowProperties properties = new FilterRowProperties("test");
    properties.init();
    FilterRowCriteriaProperties filterGreater = new FilterRowCriteriaProperties("filter1");
    filterGreater.init();
    filterGreater.columnName.setValue("a");
    filterGreater.operator.setValue(ConditionsRowConstant.Operator.GREATER);
    filterGreater.value.setValue("10");
    properties.filters.addRow(filterGreater);
    FilterRowCriteriaProperties filterLess = new FilterRowCriteriaProperties("filter2");
    filterLess.init();
    filterLess.columnName.setValue("a");
    filterLess.operator.setValue(ConditionsRowConstant.Operator.LOWER);
    filterLess.value.setValue("30");
    properties.filters.addRow(filterLess);

    FilterRowDoFn function = new FilterRowDoFn(properties);
    DoFnTester<IndexedRecord, IndexedRecord> fnTester = DoFnTester.of(function);

    List<IndexedRecord> outputs = fnTester.processBundle(input_30_300_3000_Record, input_10_100_1000_Record,
            input_20_200_2000_Record);
    List<IndexedRecord> rejects = fnTester.peekOutputElements(FilterRowRuntime.rejectOutput);

    assertEquals(1, outputs.size());
    assertEquals(2, rejects.size());

}
 

/** Tests {@link DatastoreV1.Read.SplitQueryFn} when the query has a user specified limit. */
@Test
public void testSplitQueryFnWithQueryLimit() throws Exception {
  Query queryWithLimit = QUERY.toBuilder().setLimit(Int32Value.newBuilder().setValue(1)).build();

  SplitQueryFn splitQueryFn = new SplitQueryFn(V_1_OPTIONS, 10, mockDatastoreFactory);
  DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn);
  doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
  List<Query> queries = doFnTester.processBundle(queryWithLimit);

  assertEquals(1, queries.size());
  verifyNoMoreInteractions(mockDatastore);
  verifyNoMoreInteractions(mockQuerySplitter);
}
 

@Test
public void test_FilterLogicalOpAny() throws Exception {
    FilterRowProperties properties = new FilterRowProperties("test");
    properties.init();
    properties.logicalOp.setValue(LogicalOpType.ANY);
    FilterRowCriteriaProperties condition1 = new FilterRowCriteriaProperties("filter1");
    condition1.init();
    condition1.columnName.setValue("a");
    condition1.operator.setValue(ConditionsRowConstant.Operator.EQUAL);
    condition1.value.setValue("10");
    properties.filters.addRow(condition1);
    FilterRowCriteriaProperties condition2 = new FilterRowCriteriaProperties("filter2");
    condition2.init();
    condition2.columnName.setValue("b");
    condition2.operator.setValue(ConditionsRowConstant.Operator.EQUAL);
    condition2.value.setValue("300");
    properties.filters.addRow(condition2);

    FilterRowDoFn function = new FilterRowDoFn(properties);
    DoFnTester<IndexedRecord, IndexedRecord> fnTester = DoFnTester.of(function);

    List<IndexedRecord> outputs = fnTester.processBundle(input_30_300_3000_Record, input_10_100_1000_Record,
            input_20_200_2000_Record);
    List<IndexedRecord> rejects = fnTester.peekOutputElements(FilterRowRuntime.rejectOutput);

    assertEquals(2, outputs.size());
    assertEquals(30, outputs.get(0).get(0));
    assertEquals(10, outputs.get(1).get(0));
    assertEquals(1, rejects.size());
    assertEquals(20, rejects.get(0).get(0));
}
 

@Test
public void simple() throws Exception {
  // Simplest schema: a table with int64 key
  ReadOnlyTransaction tx = mock(ReadOnlyTransaction.class);
  when(serviceFactory.mockDatabaseClient().readOnlyTransaction()).thenReturn(tx);

  preparePkMetadata(tx, Arrays.asList(pkMetadata("test", "key", "ASC")));
  prepareColumnMetadata(tx, Arrays.asList(columnMetadata("test", "key", "INT64")));

  SpannerConfig config =
      SpannerConfig.create()
          .withProjectId("test-project")
          .withInstanceId("test-instance")
          .withDatabaseId("test-database")
          .withServiceFactory(serviceFactory);

  DoFnTester<Void, SpannerSchema> tester = DoFnTester.of(new ReadSpannerSchema(config));
  List<SpannerSchema> schemas = tester.processBundle(Arrays.asList((Void) null));

  assertEquals(1, schemas.size());

  SpannerSchema schema = schemas.get(0);

  assertEquals(1, schema.getTables().size());

  SpannerSchema.Column column = SpannerSchema.Column.create("key", Type.int64());
  SpannerSchema.KeyPart keyPart = SpannerSchema.KeyPart.create("key", false);

  assertThat(schema.getColumns("test"), contains(column));
  assertThat(schema.getKeyParts("test"), contains(keyPart));
}
 

/**
 * Normalize a field not present in the input record will throw TalendRuntimeException.
 *
 * @throws Exception
 */
@Test(expected = TalendRuntimeException.class)
public void testNormalizeNotFoundField() throws Exception {
    NormalizeProperties properties = new NormalizeProperties("test");
    properties.init();
    properties.schemaListener.afterSchema();
    properties.columnToNormalize.setValue("b.y.f");
    properties.isList.setValue(false);

    NormalizeDoFn function = new NormalizeDoFn().withProperties(properties);
    DoFnTester<IndexedRecord, IndexedRecord> fnTester = DoFnTester.of(function);
    fnTester.processBundle(inputParentRecord);
}
 

@Test
public void test_FlatMap_DupplicateOutputAndApplyATransformation() throws Exception {

    PythonRowProperties properties = new PythonRowProperties("test");
    properties.init();
    properties.mapType.setValue(MapType.FLATMAP);

    StringBuilder sb = new StringBuilder();
    sb.append("import copy\n");
    sb.append("outputList.append(copy.deepcopy(input))\n");
    sb.append("outputList.append(input)\n"); // will be converted to inputIndexedRecord2
    sb.append("output = input\n");
    sb.append("output['a1'] = \"rootdata2\"\n");
    sb.append("output['B']['b1'] = \"subdatabefore\"\n");
    sb.append("output['B']['C']['c1'] = \"subsubdatabefore\"\n");
    sb.append("output['B']['C']['c2'] = 33\n");
    sb.append("output['B']['C']['c3'] = 55l\n");
    sb.append("output['B']['b2'] = \"subdataend\"\n");
    sb.append("outputList.append(output)\n");
    properties.pythonCode.setValue(sb.toString());
    PythonRowDoFn function = new PythonRowDoFn();
    assertEquals(ValidationResult.OK, function.initialize(null, properties));
    DoFnTester<IndexedRecord, IndexedRecord> fnTester = DoFnTester.of(function);
    List<IndexedRecord> outputs = fnTester.processBundle(inputIndexedRecord);
    assertEquals(3, outputs.size());

    GenericRecord outputRecord1 = (GenericRecord) outputs.get(0);
    GenericRecord outputRecord2 = (GenericRecord) outputs.get(1);
    GenericRecord outputRecord3 = (GenericRecord) outputs.get(2);

    compareRecords(inputIndexedRecord, outputRecord1);
    compareRecords(outputIndexedRecord, outputRecord2);
    compareRecords(outputIndexedRecord, outputRecord3);
}
 

private Map<KV<String, String>, KV<Double, Integer>> calculatorOutputAsMap(
    CallSimilarityCalculatorFactory calculatorFactory) throws Exception {
  DoFnTester<Variant, KV<KV<String, String>, KV<Double, Integer>>> fnTester =
      DoFnTester.of(new AlleleSimilarityCalculator(calculatorFactory));
  List<KV<KV<String, String>, KV<Double, Integer>>> fnOutput =
      fnTester.processBundle(variants.toArray(new Variant[] {}));
  Map<KV<String, String>, KV<Double, Integer>> fnOutputMap = newHashMap();
  for (KV<KV<String, String>, KV<Double, Integer>> kv : fnOutput) {
    fnOutputMap.put(kv.getKey(), kv.getValue());
  }
  return fnOutputMap;
}
 

@Test
public void selectSimpleElementsMultiplestime() throws Exception {
    FieldSelectorProperties properties = addSelector(null, "aOutput", "a");
    properties = addSelector(properties, "cOutput", "c");
    properties = addSelector(properties, "aSecondOutput", "a");
    properties = addSelector(properties, "bOutput", "b");

    FieldSelectorDoFn function = new FieldSelectorDoFn().withProperties(properties);
    DoFnTester<IndexedRecord, IndexedRecord> fnTester = DoFnTester.of(function);
    List<IndexedRecord> outputs = fnTester.processBundle(inputSimpleRecord, inputSimpleRecord);

    assertEquals(2, outputs.size());
    List<Field> fields = outputs.get(0).getSchema().getFields();
    assertEquals(4, fields.size());
    assertEquals("aOutput", fields.get(0).name());
    assertEquals("aaa", outputs.get(0).get(0));
    assertEquals("cOutput", fields.get(1).name());
    assertEquals("Ccc", outputs.get(0).get(1));
    assertEquals("aSecondOutput", fields.get(2).name());
    assertEquals("aaa", outputs.get(0).get(2));
    assertEquals("bOutput", fields.get(3).name());
    assertEquals("BBB", outputs.get(0).get(3));
    fields = outputs.get(1).getSchema().getFields();
    assertEquals(4, fields.size());
    assertEquals("aOutput", fields.get(0).name());
    assertEquals("aaa", outputs.get(1).get(0));
    assertEquals("cOutput", fields.get(1).name());
    assertEquals("Ccc", outputs.get(1).get(1));
    assertEquals("aSecondOutput", fields.get(2).name());
    assertEquals("aaa", outputs.get(1).get(2));
    assertEquals("bOutput", fields.get(3).name());
    assertEquals("BBB", outputs.get(1).get(3));
}
 

@Test
public void testBatchViewAsListToIsmRecordForGlobalWindow() throws Exception {
  DoFnTester<String, IsmRecord<WindowedValue<String>>> doFnTester =
      DoFnTester.of(
          new BatchViewOverrides.BatchViewAsList.ToIsmRecordForGlobalWindowDoFn<String>());

  // The order of the output elements is important relative to processing order
  assertThat(
      doFnTester.processBundle(ImmutableList.of("a", "b", "c")),
      contains(
          IsmRecord.of(ImmutableList.of(GlobalWindow.INSTANCE, 0L), valueInGlobalWindow("a")),
          IsmRecord.of(ImmutableList.of(GlobalWindow.INSTANCE, 1L), valueInGlobalWindow("b")),
          IsmRecord.of(ImmutableList.of(GlobalWindow.INSTANCE, 2L), valueInGlobalWindow("c"))));
}
 

@Test
public void testSplitReads_twoMatchedBasesDifferentOffsets() throws Exception {
  DoFnTester<Read, KV<Position, ReadBaseQuality>> splitReads = DoFnTester.of(new SplitReads());

  // matched bases with different offsets onto the reference
  Read r = Read.newBuilder()
      .setAlignment(LinearAlignment.newBuilder()
          .setPosition(com.google.genomics.v1.Position.newBuilder()
              .setReferenceName("1")
              .setPosition(123))
          .addCigar(CigarUnit.newBuilder()
              .setOperation(Operation.DELETE)
              .setOperationLength(1))
          .addCigar(CigarUnit.newBuilder()
              .setOperation(Operation.ALIGNMENT_MATCH)
              .setOperationLength(2))
          .addCigar(CigarUnit.newBuilder()
              .setOperation(Operation.INSERT)
              .setOperationLength(1))
          .addCigar(CigarUnit.newBuilder()
              .setOperation(Operation.ALIGNMENT_MATCH)
              .setOperationLength(1)))
       // 1D2M1I1M; A and G match positions 1 and 3 of the ref
      .setAlignedSequence("ACGT")
      .addAllAlignedQuality(ImmutableList.of(1, 2, 3, 4))
      .build();
  Assert.assertThat(splitReads.processBundle(r), CoreMatchers.hasItems(KV.of(Position.newBuilder()
          .setReferenceName("1")
          .setPosition(124L)
          .build(),
          new ReadBaseQuality("A", 1)),
      KV.of(Position.newBuilder()
          .setReferenceName("1")
          .setPosition(125L)
          .build(),
          new ReadBaseQuality("C", 2)),
      KV.of(Position.newBuilder()
          .setReferenceName("1")
          .setPosition(126L)
          .build(),
          new ReadBaseQuality("T", 4))));
}
 

void testWriteWithMaxBatchSizeBytes() throws Exception {
  Write write =
      ElasticsearchIO.write()
          .withConnectionConfiguration(connectionConfiguration)
          .withMaxBatchSizeBytes(BATCH_SIZE_BYTES);
  // write bundles size is the runner decision, we cannot force a bundle size,
  // so we test the Writer as a DoFn outside of a runner.
  try (DoFnTester<String, Void> fnTester = DoFnTester.of(new Write.WriteFn(write))) {
    List<String> input =
        ElasticsearchIOTestUtils.createDocuments(
            numDocs, ElasticsearchIOTestUtils.InjectionMode.DO_NOT_INJECT_INVALID_DOCS);
    long numDocsProcessed = 0;
    long sizeProcessed = 0;
    long numDocsInserted = 0;
    long batchInserted = 0;
    for (String document : input) {
      fnTester.processElement(document);
      numDocsProcessed++;
      sizeProcessed += document.getBytes(StandardCharsets.UTF_8).length;
      // test every 40 docs to avoid overloading ES
      if ((numDocsProcessed % 40) == 0) {
        // force the index to upgrade after inserting for the inserted docs
        // to be searchable immediately
        long currentNumDocs =
            refreshIndexAndGetCurrentNumDocs(connectionConfiguration, restClient);
        if (sizeProcessed / BATCH_SIZE_BYTES > batchInserted) {
          /* bundle end */
          assertThat(
              "we have passed a bundle size, we should have inserted some documents",
              currentNumDocs,
              greaterThan(numDocsInserted));
          numDocsInserted = currentNumDocs;
          batchInserted = (sizeProcessed / BATCH_SIZE_BYTES);
        } else {
          /* not bundle end */
          assertEquals(
              "we are not at the end of a bundle, we should have inserted no more documents",
              numDocsInserted,
              currentNumDocs);
        }
      }
    }
  }
}