Search in sources :

Example 1 with AbstractIterator

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);
}
Also used : ListDocumentsRequest(com.google.firestore.v1.ListDocumentsRequest) ListDocumentsFn(org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.ListDocumentsFn) ListDocumentsResponse(com.google.firestore.v1.ListDocumentsResponse) ListDocumentsPage(com.google.cloud.firestore.v1.FirestoreClient.ListDocumentsPage) Iterator(java.util.Iterator) AbstractIterator(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.AbstractIterator)

Example 2 with AbstractIterator

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());
}
Also used : ArgumentMatchers.any(org.mockito.ArgumentMatchers.any) RunQueryRequest(com.google.firestore.v1.RunQueryRequest) FieldReference(com.google.firestore.v1.StructuredQuery.FieldReference) ArgumentMatchers.eq(org.mockito.ArgumentMatchers.eq) Mock(org.mockito.Mock) Direction(com.google.firestore.v1.StructuredQuery.Direction) FirestoreStub(com.google.cloud.firestore.v1.stub.FirestoreStub) Operator(com.google.firestore.v1.StructuredQuery.FieldFilter.Operator) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap) Function(java.util.function.Function) AbstractIterator(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.AbstractIterator) ArgumentCaptor(org.mockito.ArgumentCaptor) Value(com.google.firestore.v1.Value) Mockito.verifyNoMoreInteractions(org.mockito.Mockito.verifyNoMoreInteractions) Objects.requireNonNull(java.util.Objects.requireNonNull) ServerStream(com.google.api.gax.rpc.ServerStream) ServerStreamingCallable(com.google.api.gax.rpc.ServerStreamingCallable) StructuredQuery(com.google.firestore.v1.StructuredQuery) Filter(com.google.firestore.v1.StructuredQuery.Filter) Document(com.google.firestore.v1.Document) Mockito.times(org.mockito.Mockito.times) Test(org.junit.Test) Mockito.doNothing(org.mockito.Mockito.doNothing) Mockito.when(org.mockito.Mockito.when) FieldFilter(com.google.firestore.v1.StructuredQuery.FieldFilter) Mockito.verify(org.mockito.Mockito.verify) RunQueryFn(org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.RunQueryFn) List(java.util.List) Cursor(com.google.firestore.v1.Cursor) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) RunQueryResponse(com.google.firestore.v1.RunQueryResponse) Order(com.google.firestore.v1.StructuredQuery.Order) CollectionSelector(com.google.firestore.v1.StructuredQuery.CollectionSelector) Collections(java.util.Collections) Assert.assertEquals(org.junit.Assert.assertEquals) RunQueryResponse(com.google.firestore.v1.RunQueryResponse) RunQueryRequest(com.google.firestore.v1.RunQueryRequest) RunQueryFn(org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.RunQueryFn)

Example 3 with AbstractIterator

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);
}
Also used : ListCollectionIdsResponse(com.google.firestore.v1.ListCollectionIdsResponse) ListCollectionIdsPage(com.google.cloud.firestore.v1.FirestoreClient.ListCollectionIdsPage) ListCollectionIdsRequest(com.google.firestore.v1.ListCollectionIdsRequest) Iterator(java.util.Iterator) AbstractIterator(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.AbstractIterator) ListCollectionIdsFn(org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.ListCollectionIdsFn)

Example 4 with AbstractIterator

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);
}
Also used : PartitionQueryPage(com.google.cloud.firestore.v1.FirestoreClient.PartitionQueryPage) PartitionQueryPair(org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.PartitionQueryPair) PartitionQueryRequest(com.google.firestore.v1.PartitionQueryRequest) Iterator(java.util.Iterator) AbstractIterator(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.AbstractIterator) PartitionQueryResponse(com.google.firestore.v1.PartitionQueryResponse) PartitionQueryFn(org.apache.beam.sdk.io.gcp.firestore.FirestoreV1ReadFn.PartitionQueryFn)

Aggregations

AbstractIterator (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.AbstractIterator)4 Iterator (java.util.Iterator)3 ServerStream (com.google.api.gax.rpc.ServerStream)1 ServerStreamingCallable (com.google.api.gax.rpc.ServerStreamingCallable)1 ListCollectionIdsPage (com.google.cloud.firestore.v1.FirestoreClient.ListCollectionIdsPage)1 ListDocumentsPage (com.google.cloud.firestore.v1.FirestoreClient.ListDocumentsPage)1 PartitionQueryPage (com.google.cloud.firestore.v1.FirestoreClient.PartitionQueryPage)1 FirestoreStub (com.google.cloud.firestore.v1.stub.FirestoreStub)1 Cursor (com.google.firestore.v1.Cursor)1 Document (com.google.firestore.v1.Document)1 ListCollectionIdsRequest (com.google.firestore.v1.ListCollectionIdsRequest)1 ListCollectionIdsResponse (com.google.firestore.v1.ListCollectionIdsResponse)1 ListDocumentsRequest (com.google.firestore.v1.ListDocumentsRequest)1 ListDocumentsResponse (com.google.firestore.v1.ListDocumentsResponse)1 PartitionQueryRequest (com.google.firestore.v1.PartitionQueryRequest)1 PartitionQueryResponse (com.google.firestore.v1.PartitionQueryResponse)1 RunQueryRequest (com.google.firestore.v1.RunQueryRequest)1 RunQueryResponse (com.google.firestore.v1.RunQueryResponse)1 StructuredQuery (com.google.firestore.v1.StructuredQuery)1 CollectionSelector (com.google.firestore.v1.StructuredQuery.CollectionSelector)1