/** Tests {@link SplitQueryFn} when number of query splits is specified. */ @Test public void testSplitQueryFnWithNumSplits() throws Exception { int numSplits = 100; when(mockQuerySplitter.getSplits( eq(QUERY), any(PartitionId.class), eq(numSplits), any(Datastore.class))) .thenReturn(splitQuery(QUERY, numSplits)); SplitQueryFn splitQueryFn = new SplitQueryFn(V_1_OPTIONS, numSplits, mockDatastoreFactory); DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn); /** * Although Datastore client is marked transient in {@link SplitQueryFn}, 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 doFn from being * serialized. */ doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE); List<Query> queries = doFnTester.processBundle(QUERY); assertEquals(queries.size(), numSplits); // Confirms that sub-queries are not equal to original when there is more than one split. for (Query subQuery : queries) { assertNotEquals(subQuery, QUERY); } verify(mockQuerySplitter, times(1)) .getSplits(eq(QUERY), any(PartitionId.class), eq(numSplits), any(Datastore.class)); verifyZeroInteractions(mockDatastore); }
/** Tests {@link SplitQueryFn} when no query splits is specified. */ @Test public void testSplitQueryFnWithoutNumSplits() throws Exception { // Force SplitQueryFn to compute the number of query splits int numSplits = 0; int expectedNumSplits = 20; long entityBytes = expectedNumSplits * DEFAULT_BUNDLE_SIZE_BYTES; // In seconds long timestamp = 1234L; RunQueryRequest latestTimestampRequest = makeRequest(makeLatestTimestampQuery(NAMESPACE), NAMESPACE); RunQueryResponse latestTimestampResponse = makeLatestTimestampResponse(timestamp); // Per Kind statistics request and response RunQueryRequest statRequest = makeRequest(makeStatKindQuery(NAMESPACE, timestamp), NAMESPACE); RunQueryResponse statResponse = makeStatKindResponse(entityBytes); when(mockDatastore.runQuery(latestTimestampRequest)).thenReturn(latestTimestampResponse); when(mockDatastore.runQuery(statRequest)).thenReturn(statResponse); when(mockQuerySplitter.getSplits( eq(QUERY), any(PartitionId.class), eq(expectedNumSplits), any(Datastore.class))) .thenReturn(splitQuery(QUERY, expectedNumSplits)); SplitQueryFn splitQueryFn = new SplitQueryFn(V_1_OPTIONS, numSplits, mockDatastoreFactory); DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn); doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE); List<Query> queries = doFnTester.processBundle(QUERY); assertEquals(expectedNumSplits, queries.size()); verify(mockQuerySplitter, times(1)) .getSplits(eq(QUERY), any(PartitionId.class), eq(expectedNumSplits), any(Datastore.class)); verify(mockDatastore, times(1)).runQuery(latestTimestampRequest); verify(mockDatastore, times(1)).runQuery(statRequest); }
@Test public void cloneOnce() throws Exception { try (DoFnTester<Long, String> tester = DoFnTester.of(new CountBundleCallsFn())) { tester.setCloningBehavior(DoFnTester.CloningBehavior.CLONE_ONCE); assertThat(tester.processBundle(1L, 2L, 3L), contains("1/0", "1/0", "1/0")); assertThat(tester.processBundle(4L, 5L), contains("2/1", "2/1")); assertThat(tester.processBundle(6L), contains("3/2")); } }
@Test public void clonePerBundle() throws Exception { try (DoFnTester<Long, String> tester = DoFnTester.of(new CountBundleCallsFn())) { tester.setCloningBehavior(DoFnTester.CloningBehavior.CLONE_PER_BUNDLE); assertThat(tester.processBundle(1L, 2L, 3L), contains("1/0", "1/0", "1/0")); assertThat(tester.processBundle(4L, 5L), contains("1/0", "1/0")); assertThat(tester.processBundle(6L), contains("1/0")); } }
@Test public void doNotClone() throws Exception { final AtomicInteger numSetupCalls = new AtomicInteger(); final AtomicInteger numTeardownCalls = new AtomicInteger(); DoFn<Long, String> fn = new DoFn<Long, String>() { @ProcessElement public void process(ProcessContext context) {} @Setup public void setup() { numSetupCalls.addAndGet(1); } @Teardown public void teardown() { numTeardownCalls.addAndGet(1); } }; try (DoFnTester<Long, String> tester = DoFnTester.of(fn)) { tester.setCloningBehavior(DoFnTester.CloningBehavior.DO_NOT_CLONE); tester.processBundle(1L, 2L, 3L); tester.processBundle(4L, 5L); tester.processBundle(6L); } assertEquals(1, numSetupCalls.get()); assertEquals(1, numTeardownCalls.get()); }
@Test public void testSupportsFinishBundleOutput() throws Exception { for (DoFnTester.CloningBehavior cloning : DoFnTester.CloningBehavior.values()) { try (DoFnTester<Integer, Integer> tester = DoFnTester.of(new BundleCounterDoFn())) { tester.setCloningBehavior(cloning); assertThat(tester.processBundle(1, 2, 3, 4), contains(4)); assertThat(tester.processBundle(5, 6, 7), contains(3)); assertThat(tester.processBundle(8, 9), contains(2)); } } }
/** 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); }
/** 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()); }
/** 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); }
/** 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 < DatastoreV1.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); }
@Test public void processBundle() throws Exception { for (DoFnTester.CloningBehavior cloning : DoFnTester.CloningBehavior.values()) { try (DoFnTester<Long, String> tester = DoFnTester.of(new CounterDoFn())) { tester.setCloningBehavior(cloning); // processBundle() returns all the output like takeOutputElements(). assertThat(tester.processBundle(1L, 2L, 3L, 4L), hasItems("1", "2", "3", "4")); // peek now returns nothing. assertTrue(tester.peekOutputElements().isEmpty()); } } }
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; } }
@Test public void processElement() throws Exception { for (DoFnTester.CloningBehavior cloning : DoFnTester.CloningBehavior.values()) { try (DoFnTester<Long, String> tester = DoFnTester.of(new CounterDoFn())) { tester.setCloningBehavior(cloning); tester.processElement(1L); List<String> take = tester.takeOutputElements(); assertThat(take, hasItems("1")); // Following takeOutputElements(), neither takeOutputElements() // nor peekOutputElements() return anything. assertTrue(tester.takeOutputElements().isEmpty()); assertTrue(tester.peekOutputElements().isEmpty()); } } }
@Test public void processMultipleBundles() throws Exception { for (DoFnTester.CloningBehavior cloning : DoFnTester.CloningBehavior.values()) { try (DoFnTester<Long, String> tester = DoFnTester.of(new CounterDoFn())) { tester.setCloningBehavior(cloning); // processBundle() returns all the output like takeOutputElements(). assertThat(tester.processBundle(1L, 2L, 3L, 4L), hasItems("1", "2", "3", "4")); assertThat(tester.processBundle(5L, 6L, 7L), hasItems("5", "6", "7")); assertThat(tester.processBundle(8L, 9L), hasItems("8", "9")); // peek now returns nothing. assertTrue(tester.peekOutputElements().isEmpty()); } } }
this.tester.setCloningBehavior(DoFnTester.CloningBehavior.DO_NOT_CLONE); this.tester.startBundle(); timerInternals.advanceProcessingTime(currentProcessingTime);
for (DoFnTester.CloningBehavior cloning : DoFnTester.CloningBehavior.values()) { try (DoFnTester<Long, String> tester = DoFnTester.of(new CounterDoFn())) { tester.setCloningBehavior(cloning);