Java DoFnTester类代码示例

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包/类
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 {@link DatastoreWriterFn} with a failed request which is retried. */
@Test
public void testDatatoreWriterFnRetriesErrors() throws Exception {
  List<Mutation> mutations = new ArrayList<>();
  int numRpcs = 2;
  for (int i = 0; i < DATASTORE_BATCH_UPDATE_ENTITIES_START * numRpcs; ++i) {
    mutations.add(
        makeUpsert(Entity.newBuilder().setKey(makeKey("key" + i, i + 1)).build()).build());
  }

  CommitResponse successfulCommit = CommitResponse.getDefaultInstance();
  when(mockDatastore.commit(any(CommitRequest.class))).thenReturn(successfulCommit)
    .thenThrow(
        new DatastoreException("commit", Code.DEADLINE_EXCEEDED, "", null))
    .thenReturn(successfulCommit);

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

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包/类
/** 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包/类
@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包/类
@Test
public void runRead() throws Exception {
  SpannerIO.Read read =
      SpannerIO.read()
          .withProjectId("test")
          .withInstanceId("123")
          .withDatabaseId("aaa")
          .withTable("users")
          .withColumns("id", "name")
          .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.read("users", KeySet.all(), Arrays.asList("id", "name")))
      .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).read("users", KeySet.all(), Arrays.asList("id", "name"));
}
 

import org.apache.beam.sdk.transforms.DoFnTester; //导入依赖的package包/类
@Test
public void testBatchViewAsSingletonToIsmRecord() 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));

  assertThat(
      doFnTester.processBundle(
          ImmutableList.of(
              KV.<Integer, Iterable<KV<GlobalWindow, WindowedValue<String>>>>of(
                  0, ImmutableList.of(KV.of(GlobalWindow.INSTANCE, valueInGlobalWindow("a")))))),
      contains(IsmRecord.of(ImmutableList.of(GlobalWindow.INSTANCE), valueInGlobalWindow("a"))));
}
 

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