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

/**
 * 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_withDot() 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());
}
 

private void datastoreWriterFnTest(int numMutations) throws Exception {
  // Create the requested number of mutations.
  List<Mutation> mutations = new ArrayList<>(numMutations);
  for (int i = 0; i < numMutations; ++i) {
    mutations.add(
        makeUpsert(Entity.newBuilder().setKey(makeKey("key" + i, i + 1)).build()).build());
  }

  DatastoreWriterFn datastoreWriter =
      new DatastoreWriterFn(
          StaticValueProvider.of(PROJECT_ID), null, mockDatastoreFactory, new FakeWriteBatcher());
  DoFnTester<Mutation, Void> doFnTester = DoFnTester.of(datastoreWriter);
  doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
  doFnTester.processBundle(mutations);

  int start = 0;
  while (start < numMutations) {
    int end = Math.min(numMutations, start + DatastoreV1.DATASTORE_BATCH_UPDATE_ENTITIES_START);
    CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
    commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
    commitRequest.addAllMutations(mutations.subList(start, end));
    // Verify all the batch requests were made with the expected mutations.
    verify(mockDatastore, times(1)).commit(commitRequest.build());
    start = end;
  }
}
 

@Ignore("Parenthesis are not correctly implemented in avpath.")
@Test
public void testHierarchical_TFD2119_B5_AtLeast1SubRecordsWithId1Or2HasValueGt10() throws Exception {
    DoFnTester<IndexedRecord, IndexedRecord> fnTester = DoFnTester.of( //
            new FilterRowDoFn(addCriteria(null, //
                    ".b1{(.id == 1) || (.id == 2)) && .value > 10}", //
                    ConditionsRowConstant.Function.COUNT, //
                    ConditionsRowConstant.Operator.GREATER, //
                    "0") //
            ));

    // TODO: for the moment, this just throws an exception.
    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 ((double) subrecord.get(2) > 10 && (id == 1 || id == 2))
                atLeastOne = true;
        }
        assertThat(main.toString(), atLeastOne, is(true));
    }
    assertThat(output, hasSize(57));
}
 

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

    // Looks like this is not an exception -- it considers .b1[99] to be the last record.
    List<IndexedRecord> output = fnTester.processBundle(inputB);
    for (IndexedRecord main : output) {
        List<IndexedRecord> subrecords = getSubrecords(main);
        assertThat(main.toString(), subrecords.get(subrecords.size() - 1).get(0), is((Object) 1));
    }
    assertThat(output, hasSize(114));
}
 

@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));
}
 

/** Tests that {@link ReadFn} retries after an error. */
@Test
public void testReadFnRetriesErrors() throws Exception {
  // An empty query to read entities.
  Query query = Query.newBuilder().setLimit(Int32Value.newBuilder().setValue(1)).build();

  // Use mockResponseForQuery to generate results.
  when(mockDatastore.runQuery(any(RunQueryRequest.class)))
      .thenThrow(new DatastoreException("RunQuery", Code.DEADLINE_EXCEEDED, "", null))
      .thenAnswer(
          invocationOnMock -> {
            Query q = ((RunQueryRequest) invocationOnMock.getArguments()[0]).getQuery();
            return mockResponseForQuery(q);
          });

  ReadFn readFn = new ReadFn(V_1_OPTIONS, mockDatastoreFactory);
  DoFnTester<Query, Entity> doFnTester = DoFnTester.of(readFn);
  doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
  doFnTester.processBundle(query);
}
 

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

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

@Test
public void test_Map_GenerateSchemaFromScratch() throws Exception {

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

    StringBuilder sb = new StringBuilder();
    sb.append("output['a1'] = \"rootdata2\"\n");
    sb.append("output['B'] = json.loads(\"{}\", object_pairs_hook=collections.OrderedDict)\n");
    sb.append("output['B']['b1'] = \"subdatabefore\"\n");
    sb.append("output['B']['C'] = json.loads(\"{}\", object_pairs_hook=collections.OrderedDict)\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);
}
 

@Test
public void testHierarchicalFirstRecordValue() throws Exception {
    FieldSelectorProperties properties = addSelector(null, "value", ".b1[0].value");
    DoFnTester<IndexedRecord, IndexedRecord> fnTester = DoFnTester.of( //
            new FieldSelectorDoFn().withProperties(properties));

    List<IndexedRecord> output = fnTester.processBundle(inputB);
    for (IndexedRecord main : output) {
        List<Field> fields = main.getSchema().getFields();
        assertEquals(1, fields.size());
        assertEquals("value", fields.get(0).name());
    }
}
 

private void checkSimpleInputInvalidOutput(DoFnTester<IndexedRecord, IndexedRecord> fnTester) throws Exception {
    List<IndexedRecord> outputs = fnTester.processBundle(inputSimpleRecord);
    assertEquals(0, outputs.size());
    List<IndexedRecord> rejects = fnTester.peekOutputElements(FilterRowRuntime.rejectOutput);
    assertEquals(1, rejects.size());
    assertEquals("aaa", rejects.get(0).get(0));
    assertEquals("BBB", rejects.get(0).get(1));
    assertEquals("Ccc", rejects.get(0).get(2));
}
 

private void checkSimpleInputValidOutput(DoFnTester<IndexedRecord, IndexedRecord> fnTester) throws Exception {
    List<IndexedRecord> outputs = fnTester.processBundle(inputSimpleRecord);
    assertEquals(1, outputs.size());
    assertEquals("aaa", outputs.get(0).get(0));
    assertEquals("BBB", outputs.get(0).get(1));
    assertEquals("Ccc", outputs.get(0).get(2));
    List<IndexedRecord> rejects = fnTester.peekOutputElements(FilterRowRuntime.rejectOutput);
    assertEquals(0, rejects.size());
}
 

@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());

}
 

@Test
public void testSharedMinorAllSimilarityFn() throws Exception {

  CallSimilarityCalculatorFactory fac = new SharedMinorAllelesCalculatorFactory();
  DoFnTester<Variant, KV<KV<String, String>, KV<Double, Integer>>> simFn =
      DoFnTester.of(new AlleleSimilarityCalculator(fac));

  List<KV<KV<String, String>, KV<Double, Integer>>> outputSnp1 = simFn.processBundle(snp1);
  assertEquals(6, outputSnp1.size());
  assertThat(outputSnp1, CoreMatchers.hasItems(
      KV.of(KV.of("het-alt sample", "ref-nocall sample"), KV.of(0.0, 1)),
      KV.of(KV.of("het-alt sample", "hom-ref sample"), KV.of(0.0, 1)),
      KV.of(KV.of("het-alt sample", "hom-alt sample"), KV.of(1.0, 1)),
      KV.of(KV.of("hom-ref sample", "ref-nocall sample"), KV.of(0.0, 1)),
      KV.of(KV.of("hom-alt sample", "ref-nocall sample"), KV.of(0.0, 1)),
      KV.of(KV.of("hom-alt sample", "hom-ref sample"), KV.of(0.0, 1))));

  List<KV<KV<String, String>, KV<Double, Integer>>> outputSnp2 = simFn.processBundle(snp2);
  assertEquals(3, outputSnp2.size());
  assertThat(
      outputSnp2,
      CoreMatchers.hasItems(KV.of(KV.of("het-alt sample", "hom-alt sample"), KV.of(0.0, 1)),
          KV.of(KV.of("het-alt sample", "ref-nocall sample"), KV.of(1.0, 1)),
          KV.of(KV.of("hom-alt sample", "ref-nocall sample"), KV.of(0.0, 1))));

  Variant[] input = new Variant[] {snp1, snp2};
  List<KV<KV<String, String>, KV<Double, Integer>>> outputBoth = simFn.processBundle(input);
  assertEquals(6, outputBoth.size());
  assertThat(outputBoth, CoreMatchers.hasItems(
      KV.of(KV.of("het-alt sample", "ref-nocall sample"), KV.of(1.0, 2)),
      KV.of(KV.of("het-alt sample", "hom-ref sample"), KV.of(0.0, 1)),
      KV.of(KV.of("het-alt sample", "hom-alt sample"), KV.of(1.0, 2)),
      KV.of(KV.of("hom-ref sample", "ref-nocall sample"), KV.of(0.0, 1)),
      KV.of(KV.of("hom-alt sample", "ref-nocall sample"), KV.of(0.0, 2)),
      KV.of(KV.of("hom-alt sample", "hom-ref sample"), KV.of(0.0, 1))));

}
 

@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));
}
 

/** Helper function to run a test reading from a {@link ReadFn}. */
private void readFnTest(int numEntities) throws Exception {
  // An empty query to read entities.
  Query query =
      Query.newBuilder().setLimit(Int32Value.newBuilder().setValue(numEntities)).build();

  // Use mockResponseForQuery to generate results.
  when(mockDatastore.runQuery(any(RunQueryRequest.class)))
      .thenAnswer(
          invocationOnMock -> {
            Query q = ((RunQueryRequest) invocationOnMock.getArguments()[0]).getQuery();
            return mockResponseForQuery(q);
          });

  ReadFn readFn = new ReadFn(V_1_OPTIONS, mockDatastoreFactory);
  DoFnTester<Query, Entity> doFnTester = DoFnTester.of(readFn);
  /**
   * Although Datastore client is marked transient in {@link ReadFn}, when injected through mock
   * factory using a when clause for unit testing purposes, it is not serializable because it
   * doesn't have a no-arg constructor. Thus disabling the cloning to prevent the test object from
   * being serialized.
   */
  doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
  List<Entity> entities = doFnTester.processBundle(query);

  int expectedNumCallsToRunQuery = (int) Math.ceil((double) numEntities / QUERY_BATCH_LIMIT);
  verify(mockDatastore, times(expectedNumCallsToRunQuery)).runQuery(any(RunQueryRequest.class));
  // Validate the number of results.
  assertEquals(numEntities, entities.size());
}
 

/**
 * Input parent record: {@link NormalizeDoFnTest#inputParentRecord}
 *
 * Normalize simple field: `b.y.d.j.l`
 *
 * Throw an exception: the element l is inside a loop.
 *
 * @throws Exception
 */
@Test(expected = TalendRuntimeException.class)
public void testNormalizeArrayFields_bydjl() throws Exception {
    NormalizeProperties properties = new NormalizeProperties("test");
    properties.init();
    properties.schemaListener.afterSchema();
    // Normalize `b.y.d.j` array field
    properties.columnToNormalize.setValue("b.y.d.j.l");

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

/**
 * Input parent record: {@link NormalizeDoFnTest#inputParentRecord}
 *
 * Normalize simple field: `b.y.d.k.t`
 *
 * Throw an exception: the element t does not exist
 *
 * @throws Exception
 */
@Test(expected = TalendRuntimeException.class)
public void testNormalizeSimpleFields_bydkt() throws Exception {
    NormalizeProperties properties = new NormalizeProperties("test");
    properties.init();
    properties.schemaListener.afterSchema();
    // Normalize `b.y.d.k.t` simple field
    properties.columnToNormalize.setValue("b.y.d.k.t");

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

@Test
public void test_utf8() throws Exception {
    PythonRowProperties properties = new PythonRowProperties("test");
    properties.init();
    properties.mapType.setValue(MapType.MAP);
    properties.pythonCode.setValue("output['a1'] = input['a1']");
    PythonRowDoFn function = new PythonRowDoFn();
    assertEquals(ValidationResult.OK, function.initialize(null, properties));
    DoFnTester<IndexedRecord, IndexedRecord> fnTester = DoFnTester.of(function);
    List<IndexedRecord> outputs = fnTester.processBundle(GenericDataRecordHelper.createRecord(new Object[] { utf8Sample }));
    assertEquals(utf8Sample, outputs.get(0).get(0));
}
 

@Test(expected = TalendRuntimeException.class)
public void selectInvalidElements() throws Exception {
    FieldSelectorProperties properties = addSelector(null, "a b", "a");

    FieldSelectorDoFn function = new FieldSelectorDoFn().withProperties(properties);
    DoFnTester<IndexedRecord, IndexedRecord> fnTester = DoFnTester.of(function);
    fnTester.processBundle(inputSimpleRecord);
    // throw exception
}
 

@Test
public void testToIsmRecordForMapLikeDoFnWithoutUniqueKeysThrowsException() throws Exception {
  TupleTag<KV<Integer, KV<IntervalWindow, Long>>> outputForSizeTag = new TupleTag<>();
  TupleTag<KV<Integer, KV<IntervalWindow, Long>>> outputForEntrySetTag = new TupleTag<>();

  Coder<Long> keyCoder = VarLongCoder.of();
  Coder<IntervalWindow> windowCoder = IntervalWindow.getCoder();

  IsmRecordCoder<WindowedValue<Long>> ismCoder =
      IsmRecordCoder.of(
          1,
          2,
          ImmutableList.of(
              MetadataKeyCoder.of(keyCoder), IntervalWindow.getCoder(), BigEndianLongCoder.of()),
          FullWindowedValueCoder.of(VarLongCoder.of(), windowCoder));

  DoFnTester<
          KV<Integer, Iterable<KV<KV<Long, IntervalWindow>, WindowedValue<Long>>>>,
          IsmRecord<WindowedValue<Long>>>
      doFnTester =
          DoFnTester.of(
              new BatchViewOverrides.BatchViewAsMultimap.ToIsmRecordForMapLikeDoFn<>(
                  outputForSizeTag,
                  outputForEntrySetTag,
                  windowCoder,
                  keyCoder,
                  ismCoder,
                  true /* unique keys */));

  IntervalWindow windowA = new IntervalWindow(new Instant(0), new Instant(10));

  Iterable<KV<Integer, Iterable<KV<KV<Long, IntervalWindow>, WindowedValue<Long>>>>>
      inputElements =
          ImmutableList.of(
              KV.of(
                  1,
                  (Iterable<KV<KV<Long, IntervalWindow>, WindowedValue<Long>>>)
                      ImmutableList.of(
                          KV.of(
                              KV.of(1L, windowA),
                              WindowedValue.of(
                                  110L, new Instant(1), windowA, PaneInfo.NO_FIRING)),
                          // same window same key as to previous
                          KV.of(
                              KV.of(1L, windowA),
                              WindowedValue.of(
                                  111L, new Instant(2), windowA, PaneInfo.NO_FIRING)))));

  thrown.expect(IllegalStateException.class);
  thrown.expectMessage("Unique keys are expected but found key");
  doFnTester.processBundle(inputElements);
}