Java DoFnTester.of方法代码示例

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
@Test
public void testDatastoreToGcs_EntityToJson_noTransform() throws Exception {
  DoFnTester<Entity, String> fnTester = DoFnTester.of(EntityToJson.newBuilder()
      .setJsTransformPath(StaticValueProvider.of(null))
      .setJsTransformFunctionName(StaticValueProvider.of(null))
      .build());

  Builder entityBuilder = Entity.newBuilder();
  JsonFormat.parser().usingTypeRegistry(
      TypeRegistry.newBuilder()
          .add(Entity.getDescriptor())
          .build())
      .merge(mEntityJson, entityBuilder);

  Entity entity = entityBuilder.build();
  List<String> entityJsonOutputs = fnTester.processBundle(entity);
  Assert.assertEquals(mEntityJson, entityJsonOutputs.get(0));
}
 

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
@Test
public void testDatastoreToGcs_EntityToJson_withTransform() throws Exception {
  DoFnTester<Entity, String> fnTester = DoFnTester.of(EntityToJson.newBuilder()
      .setJsTransformPath(StaticValueProvider.of(jsTransformPath))
      .setJsTransformFunctionName(StaticValueProvider.of("transform"))
      .build());

  Builder entityBuilder = Entity.newBuilder();
  JsonFormat.parser().usingTypeRegistry(
      TypeRegistry.newBuilder()
          .add(Entity.getDescriptor())
          .build())
      .merge(mEntityJson, entityBuilder);

  Entity entity = entityBuilder.build();
  List<String> entityJsonOutputs = fnTester.processBundle(entity);
  Assert.assertEquals(mTransformedEntityJson, entityJsonOutputs.get(0));
}
 

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
@Test
public void testGcsToDatastore_EntityToJson_noTransform() throws Exception {
  DoFnTester<String, Entity> fnTester = DoFnTester.of(JsonToEntity.newBuilder()
      .setJsTransformPath(StaticValueProvider.of(null))
      .setJsTransformFunctionName(StaticValueProvider.of(null))
      .build());
  List<Entity> output = fnTester.processBundle(mEntityJson);
  Entity outputEntity = output.get(0);

  Printer printer = JsonFormat.printer()
      .omittingInsignificantWhitespace()
      .usingTypeRegistry(
          TypeRegistry.newBuilder()
              .add(Entity.getDescriptor())
              .build());
  Assert.assertEquals(mEntityJson, printer.print(outputEntity));
}
 

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
@Test
public void testDatastoreToBq_EntityToTableRow_notransform() throws Exception, IOException {
  DoFnTester<Entity, TableRow> fnTester = DoFnTester.of(EntityToTableRow.newBuilder()
      .setStrictCast(StaticValueProvider.of(true))
      .setTableSchemaJson(StaticValueProvider.of(mTableSchemaJson))
      .setJsTransformFunctionName(StaticValueProvider.of(null))
      .setJsTransformPath(StaticValueProvider.of(null))
  .build());

  Builder entityBuilder = Entity.newBuilder();
  JsonFormat.parser().usingTypeRegistry(
      TypeRegistry.newBuilder()
          .add(Entity.getDescriptor())
          .build())
      .merge(mEntityJson, entityBuilder);

  Entity entity = entityBuilder.build();
  List<TableRow> tableRows = fnTester.processBundle(entity);
  TableRow tr = tableRows.get(0);

  Assert.assertEquals(1, tableRows.size());
  Assert.assertEquals("key(Drawing, '31ce830e-91d0-405e-855a-abe416cadc1f')", tr.get("__key__"));
  Assert.assertEquals("79a1d9d9-e255-427a-9b09-f45157e97790", tr.get("canvasId"));
}
 

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
@Test
public void testApply() {
  ObjectMapper mapperObj = new ObjectMapper();
  Map map = new HashMap();
  map.put("foo", 1);
  map.put("bar", Arrays.asList(1, 2, 3));
  Map submap = new HashMap();
  submap.put("moge", 1);
  map.put("hoge", submap);
  String jsonStr = new Gson().toJson(map);

  PubsubMessage2TableRowFn fn = new PubsubMessage2TableRowFn();
  DoFnTester<String, TableRow> fnTester = DoFnTester.of(fn);

  try {
    List<TableRow> result = fnTester.processBundle(jsonStr);
    String msg = result.get(0).get(PubsubMessage2TableRowFn.getMessageColumnName()).toString();
    assertEquals(jsonStr, msg);
  } catch (Exception e) {
    e.printStackTrace();
  }
}
 

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
@Test
public void testUserAgentAnalysisDoFn() throws Exception {
    DoFn<TestRecord, TestRecord> fn = new TestDoFn();

    DoFnTester<TestRecord, TestRecord> fnTester = DoFnTester.of(fn);

    // Process a bundle containing a single input element:
    TestRecord testInput = new TestRecord("Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) " +
        "Chrome/48.0.2564.82 Safari/537.36");
    List<TestRecord> testOutputs = fnTester.processBundle(testInput);

    assertEquals(1, testOutputs.size());
    TestRecord record = testOutputs.get(0);
    assertEquals("Desktop", record.deviceClass);
    assertEquals("Chrome 48.0.2564.82", record.agentNameVersion);
    assertNull(record.shouldRemainNull);
}
 

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
@Test
@SuppressWarnings({"rawtypes", "unchecked"})
public void testFormatCounts() throws Exception {
  DoFnTester<KV<Integer, Long>, TableRow> formatCountsFn =
      DoFnTester.of(new FormatCountsFn());
  KV empty[] = {};
  List<TableRow> results = formatCountsFn.processBundle(empty);
  Assert.assertTrue(results.size() == 0);
  KV input[] = { KV.of(3, 0L),
                 KV.of(4, Long.MAX_VALUE),
                 KV.of(5, Long.MIN_VALUE) };
  results = formatCountsFn.processBundle(input);
  Assert.assertEquals(results.size(), 3);
  Assert.assertEquals(results.get(0).get("month"), 3);
  Assert.assertEquals(results.get(0).get("tornado_count"), 0L);
  Assert.assertEquals(results.get(1).get("month"), 4);
  Assert.assertEquals(results.get(1).get("tornado_count"), Long.MAX_VALUE);
  Assert.assertEquals(results.get(2).get("month"), 5);
  Assert.assertEquals(results.get(2).get("tornado_count"), Long.MIN_VALUE);
}
 

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
@Test
public void runQuery() throws Exception {
  SpannerIO.Read read =
      SpannerIO.read()
          .withProjectId("test")
          .withInstanceId("123")
          .withDatabaseId("aaa")
          .withQuery("SELECT * FROM users")
          .withServiceFactory(serviceFactory);

  NaiveSpannerReadFn readFn = new NaiveSpannerReadFn(read.getSpannerConfig());
  DoFnTester<ReadOperation, Struct> fnTester = DoFnTester.of(readFn);

  when(serviceFactory.mockDatabaseClient().readOnlyTransaction(any(TimestampBound.class)))
      .thenReturn(mockTx);
  when(mockTx.executeQuery(any(Statement.class)))
      .thenReturn(ResultSets.forRows(FAKE_TYPE, FAKE_ROWS));

  List<Struct> result = fnTester.processBundle(read.getReadOperation());
  assertThat(result, Matchers.containsInAnyOrder(FAKE_ROWS.toArray()));

  verify(serviceFactory.mockDatabaseClient()).readOnlyTransaction(TimestampBound
      .strong());
  verify(mockTx).executeQuery(Statement.of("SELECT * FROM users"));
}
 

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
/**
 * 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(queries.size(), 1);
  verifyNoMoreInteractions(mockDatastore);
  verifyNoMoreInteractions(mockQuerySplitter);
}
 

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
@Test
public void testHierarchical_TFD2119_A2_Subrecord() throws Exception {
    DoFnTester<Object, IndexedRecord> fnTester = DoFnTester.of( //
            new FilterRowDoFn().withProperties(addCriteria(null, //
                    ".a1.id", //
                    null, //
                    ConditionsRowConstant.Operator.EQUAL, //
                    "1") //

            ).withOutputSchema(true));

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

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
@Test
public void testHierarchical_TFD2119_A3_Subsubrecord() throws Exception {
    DoFnTester<Object, IndexedRecord> fnTester = DoFnTester.of( //
            new FilterRowDoFn().withProperties(addCriteria(null, //
                    ".a1.a2.id", //
                    null, //
                    ConditionsRowConstant.Operator.EQUAL, //
                    "1") //

            ).withOutputSchema(true));

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

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
@Test
public void test_FlatMap_MultipleInputs() 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(input)\n");
    sb.append("outputList.append(input)\n");
    sb.append("outputList.append(copy.deepcopy(input))\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, inputIndexedRecord, inputIndexedRecord);
    assertEquals(9, outputs.size());

    for (int i = 0; i < 9; i++) {
        GenericRecord outputRecord = (GenericRecord) outputs.get(i);
        compareRecords(inputIndexedRecord, outputRecord);
    }
}
 

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
/**
 * 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());
}
 

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
@Test
public void testHierarchical_TFD2119_B6_AllSubRecordsWithId1Or2HasValueGt10() throws Exception {
    DoFnTester<Object, IndexedRecord> fnTester = DoFnTester.of( //
            new FilterRowDoFn().withProperties(addCriteria(null, //
                    ".b1{.id == 1 || .id == 2}.value", //
                    null, //
                    ConditionsRowConstant.Operator.GREATER, //
                    "10") //

            ).withOutputSchema(true));

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

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
@Test
public void testHierarchical_TFD2119_B4_LastRecordValueGt10() throws Exception {
    DoFnTester<Object, IndexedRecord> fnTester = DoFnTester.of( //
            new FilterRowDoFn().withProperties(addCriteria(null, //
                    ".b1[-1].value", //
                    null, //
                    ConditionsRowConstant.Operator.GREATER, //
                    "10") //

            ).withOutputSchema(true));

    List<IndexedRecord> output = fnTester.processBundle(inputB);
    for (IndexedRecord main : output) {
        List<IndexedRecord> subrecords = getSubrecords(main);
        assertThat(main.toString(), (double) subrecords.get(subrecords.size() - 1).get(2), greaterThan(10d));
    }
    assertThat(output, hasSize(145));
}