Examples with BatchGetDocumentsFn org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.BatchGetDocumentsFn used on opensource projects

Example 1 with BatchGetDocumentsFn

use of org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.BatchGetDocumentsFn in project beam by apache.

the class FirestoreV1FnBatchGetDocumentsTest method resumeFromLastReadValue.

@Override
public void resumeFromLastReadValue() throws Exception {
    final BatchGetDocumentsResponse response1 = newMissing(1);
    final BatchGetDocumentsResponse response2 = newFound(2);
    final BatchGetDocumentsResponse response3 = newMissing(3);
    final BatchGetDocumentsResponse response4 = newFound(4);
    final BatchGetDocumentsRequest request1 = BatchGetDocumentsRequest.newBuilder().setDatabase(String.format("projects/%s/databases/(default)", projectId)).addDocuments(response1.getMissing()).addDocuments(response2.getFound().getName()).addDocuments(response3.getMissing()).addDocuments(response4.getFound().getName()).build();
    BatchGetDocumentsRequest request2 = BatchGetDocumentsRequest.newBuilder().setDatabase(String.format("projects/%s/databases/(default)", projectId)).addDocuments(response3.getMissing()).addDocuments(response4.getFound().getName()).build();
    BatchGetDocumentsRequest request3 = BatchGetDocumentsRequest.newBuilder().setDatabase(String.format("projects/%s/databases/(default)", projectId)).addDocuments(response4.getFound().getName()).build();
    when(responseStream1.iterator()).thenReturn(new AbstractIterator<BatchGetDocumentsResponse>() {

        private int counter = 10;

        @Override
        protected BatchGetDocumentsResponse computeNext() {
            int count = counter++;
            if (count == 10) {
                return response1;
            } else if (count == 11) {
                return response2;
            } else {
                throw RETRYABLE_ERROR;
            }
        }
    });
    when(responseStream2.iterator()).thenReturn(new AbstractIterator<BatchGetDocumentsResponse>() {

        private int counter = 20;

        @Override
        protected BatchGetDocumentsResponse computeNext() {
            int count = counter++;
            if (count == 20) {
                return response3;
            } else {
                throw RETRYABLE_ERROR;
            }
        }
    });
    when(responseStream3.iterator()).thenReturn(ImmutableList.of(response4).iterator());
    doNothing().when(attempt).checkCanRetry(any(), eq(RETRYABLE_ERROR));
    when(callable.call(request1)).thenReturn(responseStream1);
    when(callable.call(request2)).thenReturn(responseStream2);
    when(callable.call(request3)).thenReturn(responseStream3);
    when(stub.batchGetDocumentsCallable()).thenReturn(callable);
    when(ff.getFirestoreStub(any())).thenReturn(stub);
    when(ff.getRpcQos(any())).thenReturn(rpcQos);
    when(rpcQos.newReadAttempt(any())).thenReturn(attempt);
    when(attempt.awaitSafeToProceed(any())).thenReturn(true);
    ArgumentCaptor<BatchGetDocumentsResponse> responsesCaptor = ArgumentCaptor.forClass(BatchGetDocumentsResponse.class);
    doNothing().when(processContext).output(responsesCaptor.capture());
    when(processContext.element()).thenReturn(request1);
    BatchGetDocumentsFn fn = new BatchGetDocumentsFn(clock, ff, rpcQosOptions);
    runFunction(fn);
    List<BatchGetDocumentsResponse> expectedResponses = ImmutableList.of(response1, response2, response3, response4);
    List<BatchGetDocumentsResponse> actualResponses = responsesCaptor.getAllValues();
    assertEquals(expectedResponses, actualResponses);
    verify(callable, times(1)).call(request1);
    verify(callable, times(1)).call(request2);
    verify(attempt, times(4)).recordStreamValue(any());
}