use of com.google.firestore.v1.BatchGetDocumentsResponse 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());
}
use of com.google.firestore.v1.BatchGetDocumentsResponse in project beam by apache.
the class FirestoreV1FnBatchGetDocumentsTest method endToEnd.
@Test
public void endToEnd() throws Exception {
final BatchGetDocumentsRequest request = BatchGetDocumentsRequest.newBuilder().setDatabase(String.format("projects/%s/databases/(default)", projectId)).addDocuments("doc_1-1").addDocuments("doc_1-2").addDocuments("doc_1-3").build();
final BatchGetDocumentsResponse response1 = newFound(1);
final BatchGetDocumentsResponse response2 = newFound(2);
final BatchGetDocumentsResponse response3 = newFound(3);
List<BatchGetDocumentsResponse> responses = ImmutableList.of(response1, response2, response3);
when(responseStream1.iterator()).thenReturn(responses.iterator());
when(callable.call(request)).thenReturn(responseStream1);
when(stub.batchGetDocumentsCallable()).thenReturn(callable);
when(ff.getFirestoreStub(any())).thenReturn(stub);
when(ff.getRpcQos(any())).thenReturn(FirestoreStatefulComponentFactory.INSTANCE.getRpcQos(rpcQosOptions));
ArgumentCaptor<BatchGetDocumentsResponse> responsesCaptor = ArgumentCaptor.forClass(BatchGetDocumentsResponse.class);
doNothing().when(processContext).output(responsesCaptor.capture());
when(processContext.element()).thenReturn(request);
runFunction(getFn(clock, ff, rpcQosOptions));
List<BatchGetDocumentsResponse> allValues = responsesCaptor.getAllValues();
assertEquals(responses, allValues);
}
Aggregations