Java DoFnTester.processBundle方法代码示例

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
@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 testBatchViewAsSingletonToIsmRecordWithMultipleValuesThrowsException()
    throws Exception {
  DoFnTester<
          KV<Integer, Iterable<KV<GlobalWindow, WindowedValue<String>>>>,
          IsmRecord<WindowedValue<String>>>
      doFnTester =
          DoFnTester.of(
              new BatchViewOverrides.BatchViewAsSingleton.IsmRecordForSingularValuePerWindowDoFn<
                  String, GlobalWindow>(GlobalWindow.Coder.INSTANCE));

  thrown.expect(IllegalStateException.class);
  thrown.expectMessage("found for singleton within window");
  doFnTester.processBundle(
      ImmutableList.of(
          KV.<Integer, Iterable<KV<GlobalWindow, WindowedValue<String>>>>of(
              0,
              ImmutableList.of(
                  KV.of(GlobalWindow.INSTANCE, valueInGlobalWindow("a")),
                  KV.of(GlobalWindow.INSTANCE, valueInGlobalWindow("b"))))));
}
 

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

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
/** 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(new Answer<RunQueryResponse>() {
        @Override
        public RunQueryResponse answer(InvocationOnMock invocationOnMock) throws Throwable {
          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);
  List<Entity> entities = doFnTester.processBundle(query);
}
 

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
/**
 * Invalid path to normalize => throw error
 */
@Test
public void testNormalize_nothing() 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(null);

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

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() 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包/类
/**
 * 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_B1_AtLeastOneSubRecordHasValueGt10() throws Exception {
    DoFnTester<Object, IndexedRecord> fnTester = DoFnTester.of( //
            new FilterRowDoFn().withProperties(addCriteria(null, //
                    ".b1{.value > 10}", //
                    ConditionsRowConstant.Function.COUNT, //
                    ConditionsRowConstant.Operator.GREATER, //
                    "0") //

            ).withOutputSchema(true));

    List<IndexedRecord> output = fnTester.processBundle(inputB);
    for (IndexedRecord main : output) {
        boolean atLeastOne = false;
        for (IndexedRecord subrecord : getSubrecords(main)) {
            if ((double) subrecord.get(2) > 10)
                atLeastOne = true;
        }
        assertThat(main.toString(), atLeastOne, is(true));
    }
    assertThat(output, hasSize(274));
}
 

import org.apache.beam.sdk.transforms.DoFnTester; //导入方法依赖的package包/类
@Test
public void test_Map_doNothing() throws Exception {

    PythonRowProperties properties = new PythonRowProperties("test");
    properties.init();
    properties.mapType.setValue(MapType.MAP);
    properties.pythonCode.setValue("output = input");
    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(inputIndexedRecord, outputRecord);
}
 

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