use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.AbstractIterator in project beam by apache.
the class FirestoreV1FnListDocumentsTest method resumeFromLastReadValue.
@Override
public void resumeFromLastReadValue() throws Exception {
when(ff.getFirestoreStub(any())).thenReturn(stub);
when(ff.getRpcQos(any())).thenReturn(rpcQos);
when(rpcQos.newReadAttempt(any())).thenReturn(attempt);
when(attempt.awaitSafeToProceed(any())).thenReturn(true);
// First page of the response
ListDocumentsRequest request1 = ListDocumentsRequest.newBuilder().setParent(String.format("projects/%s/databases/(default)/document", projectId)).build();
ListDocumentsResponse response1 = ListDocumentsResponse.newBuilder().addDocuments(Document.newBuilder().setName("doc_1-1").putAllFields(ImmutableMap.of("foo", Value.newBuilder().setStringValue("bar").build())).build()).addDocuments(Document.newBuilder().setName("doc_1-2").putAllFields(ImmutableMap.of("foo", Value.newBuilder().setStringValue("bar").build())).build()).addDocuments(Document.newBuilder().setName("doc_1-3").putAllFields(ImmutableMap.of("foo", Value.newBuilder().setStringValue("bar").build())).build()).setNextPageToken("page2").build();
when(page1.getNextPageToken()).thenReturn(response1.getNextPageToken());
when(page1.getResponse()).thenReturn(response1);
when(page1.hasNextPage()).thenReturn(true);
when(callable.call(request1)).thenReturn(pagedResponse1);
doNothing().when(attempt).checkCanRetry(any(), eq(RETRYABLE_ERROR));
when(pagedResponse1.iteratePages()).thenAnswer(invocation -> new Iterable<ListDocumentsPage>() {
@Override
public Iterator<ListDocumentsPage> iterator() {
return new AbstractIterator<ListDocumentsPage>() {
private boolean first = true;
@Override
protected ListDocumentsPage computeNext() {
if (first) {
first = false;
return page1;
} else {
throw RETRYABLE_ERROR;
}
}
};
}
});
// Second page of the response
ListDocumentsRequest request2 = ListDocumentsRequest.newBuilder().setParent(String.format("projects/%s/databases/(default)/document", projectId)).setPageToken("page2").build();
ListDocumentsResponse response2 = ListDocumentsResponse.newBuilder().addDocuments(Document.newBuilder().setName("doc_2-1").putAllFields(ImmutableMap.of("foo", Value.newBuilder().setStringValue("bar").build())).build()).build();
when(page2.getResponse()).thenReturn(response2);
when(page2.hasNextPage()).thenReturn(false);
when(callable.call(request2)).thenReturn(pagedResponse2);
when(pagedResponse2.iteratePages()).thenReturn(ImmutableList.of(page2));
when(stub.listDocumentsPagedCallable()).thenReturn(callable);
when(ff.getFirestoreStub(any())).thenReturn(stub);
ArgumentCaptor<ListDocumentsResponse> responses = ArgumentCaptor.forClass(ListDocumentsResponse.class);
doNothing().when(processContext).output(responses.capture());
when(processContext.element()).thenReturn(request1);
ListDocumentsFn fn = new ListDocumentsFn(clock, ff, rpcQosOptions);
runFunction(fn);
List<ListDocumentsResponse> expected = newArrayList(response1, response2);
List<ListDocumentsResponse> allValues = responses.getAllValues();
assertEquals(expected, allValues);
}
use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.AbstractIterator in project beam by apache.
the class FirestoreV1FnRunQueryTest method resumeFromLastReadValue.
@Override
public void resumeFromLastReadValue() throws Exception {
TestData testData = TestData.fieldEqualsBar().setProjectId(projectId).setOrderFunction(f -> Collections.singletonList(Order.newBuilder().setDirection(Direction.ASCENDING).setField(f).build())).build();
RunQueryRequest request2 = RunQueryRequest.newBuilder().setParent(String.format("projects/%s/databases/(default)/document", projectId)).setStructuredQuery(testData.request.getStructuredQuery().toBuilder().setStartAt(Cursor.newBuilder().setBefore(false).addValues(Value.newBuilder().setStringValue("bar")))).build();
List<RunQueryResponse> responses = ImmutableList.of(testData.response1, testData.response2, testData.response3);
when(responseStream1.iterator()).thenReturn(new AbstractIterator<RunQueryResponse>() {
private int invocationCount = 1;
@Override
protected RunQueryResponse computeNext() {
int count = invocationCount++;
if (count == 1) {
return responses.get(0);
} else if (count == 2) {
return responses.get(1);
} else {
throw RETRYABLE_ERROR;
}
}
});
when(callable.call(testData.request)).thenReturn(responseStream1);
doNothing().when(attempt).checkCanRetry(any(), eq(RETRYABLE_ERROR));
when(responseStream2.iterator()).thenReturn(ImmutableList.of(responses.get(2)).iterator());
when(callable.call(request2)).thenReturn(responseStream2);
when(stub.runQueryCallable()).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<RunQueryResponse> responsesCaptor = ArgumentCaptor.forClass(RunQueryResponse.class);
doNothing().when(processContext).output(responsesCaptor.capture());
when(processContext.element()).thenReturn(testData.request);
RunQueryFn fn = new RunQueryFn(clock, ff, rpcQosOptions);
runFunction(fn);
List<RunQueryResponse> allValues = responsesCaptor.getAllValues();
assertEquals(responses, allValues);
verify(callable, times(1)).call(testData.request);
verify(callable, times(1)).call(request2);
verify(attempt, times(3)).recordStreamValue(any());
}
use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.AbstractIterator in project beam by apache.
the class FirestoreV1FnListCollectionIdsTest method resumeFromLastReadValue.
@Override
public void resumeFromLastReadValue() throws Exception {
when(ff.getFirestoreStub(any())).thenReturn(stub);
when(ff.getRpcQos(any())).thenReturn(rpcQos);
when(rpcQos.newReadAttempt(any())).thenReturn(attempt);
when(attempt.awaitSafeToProceed(any())).thenReturn(true);
// First page of the response
ListCollectionIdsRequest request1 = ListCollectionIdsRequest.newBuilder().setParent(String.format("projects/%s/databases/(default)/document", projectId)).build();
ListCollectionIdsResponse response1 = ListCollectionIdsResponse.newBuilder().addCollectionIds("col_1-1").addCollectionIds("col_1-2").addCollectionIds("col_1-3").setNextPageToken("page2").build();
when(page1.getNextPageToken()).thenReturn(response1.getNextPageToken());
when(page1.getResponse()).thenReturn(response1);
when(page1.hasNextPage()).thenReturn(true);
when(callable.call(request1)).thenReturn(pagedResponse1);
doNothing().when(attempt).checkCanRetry(any(), eq(RETRYABLE_ERROR));
when(pagedResponse1.iteratePages()).thenAnswer(invocation -> new Iterable<ListCollectionIdsPage>() {
@Override
public Iterator<ListCollectionIdsPage> iterator() {
return new AbstractIterator<ListCollectionIdsPage>() {
private boolean first = true;
@Override
protected ListCollectionIdsPage computeNext() {
if (first) {
first = false;
return page1;
} else {
throw RETRYABLE_ERROR;
}
}
};
}
});
// Second page of the response
ListCollectionIdsRequest request2 = ListCollectionIdsRequest.newBuilder().setParent(String.format("projects/%s/databases/(default)/document", projectId)).setPageToken("page2").build();
ListCollectionIdsResponse response2 = ListCollectionIdsResponse.newBuilder().addCollectionIds("col_2-1").build();
when(page2.getResponse()).thenReturn(response2);
when(page2.hasNextPage()).thenReturn(false);
when(callable.call(request2)).thenReturn(pagedResponse2);
when(pagedResponse2.iteratePages()).thenReturn(ImmutableList.of(page2));
when(stub.listCollectionIdsPagedCallable()).thenReturn(callable);
when(ff.getFirestoreStub(any())).thenReturn(stub);
ArgumentCaptor<ListCollectionIdsResponse> responses = ArgumentCaptor.forClass(ListCollectionIdsResponse.class);
doNothing().when(processContext).output(responses.capture());
when(processContext.element()).thenReturn(request1);
ListCollectionIdsFn fn = new ListCollectionIdsFn(clock, ff, rpcQosOptions);
runFunction(fn);
List<ListCollectionIdsResponse> expected = newArrayList(response1, response2);
List<ListCollectionIdsResponse> allValues = responses.getAllValues();
assertEquals(expected, allValues);
}
use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.AbstractIterator in project beam by apache.
the class FirestoreV1FnPartitionQueryTest method resumeFromLastReadValue.
@Override
public void resumeFromLastReadValue() throws Exception {
when(ff.getFirestoreStub(any())).thenReturn(stub);
when(ff.getRpcQos(any())).thenReturn(rpcQos);
when(rpcQos.newReadAttempt(any())).thenReturn(attempt);
when(attempt.awaitSafeToProceed(any())).thenReturn(true);
// First page of the response
PartitionQueryRequest request1 = PartitionQueryRequest.newBuilder().setParent(String.format("projects/%s/databases/(default)/document", projectId)).build();
PartitionQueryResponse response1 = PartitionQueryResponse.newBuilder().addPartitions(Cursor.newBuilder().addValues(Value.newBuilder().setReferenceValue("doc-100"))).addPartitions(Cursor.newBuilder().addValues(Value.newBuilder().setReferenceValue("doc-200"))).addPartitions(Cursor.newBuilder().addValues(Value.newBuilder().setReferenceValue("doc-300"))).setNextPageToken("page2").build();
when(page1.getResponse()).thenReturn(response1);
when(callable.call(request1)).thenReturn(pagedResponse1);
doNothing().when(attempt).checkCanRetry(any(), eq(RETRYABLE_ERROR));
when(pagedResponse1.iteratePages()).thenAnswer(invocation -> new Iterable<PartitionQueryPage>() {
@Override
public Iterator<PartitionQueryPage> iterator() {
return new AbstractIterator<PartitionQueryPage>() {
private boolean first = true;
@Override
protected PartitionQueryPage computeNext() {
if (first) {
first = false;
return page1;
} else {
throw RETRYABLE_ERROR;
}
}
};
}
});
// Second page of the response
PartitionQueryRequest request2 = PartitionQueryRequest.newBuilder().setParent(String.format("projects/%s/databases/(default)/document", projectId)).setPageToken("page2").build();
PartitionQueryResponse response2 = PartitionQueryResponse.newBuilder().addPartitions(Cursor.newBuilder().addValues(Value.newBuilder().setReferenceValue("doc-400"))).build();
PartitionQueryResponse expectedResponse = response1.toBuilder().clearNextPageToken().addAllPartitions(response2.getPartitionsList()).build();
when(page2.getResponse()).thenReturn(response2);
when(page2.hasNextPage()).thenReturn(false);
when(callable.call(request2)).thenReturn(pagedResponse2);
when(pagedResponse2.iteratePages()).thenReturn(ImmutableList.of(page2));
when(stub.partitionQueryPagedCallable()).thenReturn(callable);
when(ff.getFirestoreStub(any())).thenReturn(stub);
ArgumentCaptor<PartitionQueryPair> responses = ArgumentCaptor.forClass(PartitionQueryPair.class);
doNothing().when(processContext).output(responses.capture());
when(processContext.element()).thenReturn(request1);
PartitionQueryFn fn = new PartitionQueryFn(clock, ff, rpcQosOptions);
runFunction(fn);
List<PartitionQueryPair> expected = newArrayList(new PartitionQueryPair(request1, expectedResponse));
List<PartitionQueryPair> allValues = responses.getAllValues();
assertEquals(expected, allValues);
}
Aggregations