Examples with TestDataLayoutHint org.apache.beam.sdk.io.gcp.firestore.it.FirestoreTestingHelper.TestDataLayoutHint used on opensource projects

Example 1 with TestDataLayoutHint

use of org.apache.beam.sdk.io.gcp.firestore.it.FirestoreTestingHelper.TestDataLayoutHint in project beam by apache.

the class BaseFirestoreIT method listCollections.

@Test
@TestDataLayoutHint(DataLayout.Deep)
public final void listCollections() throws Exception {
    // verification and cleanup of nested collections is much slower because each document
    // requires an rpc to find its collections, instead of using the usual size, use 20
    // to keep the test quick
    List<String> collectionIds = IntStream.rangeClosed(1, 20).mapToObj(i -> helper.colId()).collect(Collectors.toList());
    ApiFutures.transform(ApiFutures.allAsList(chunkUpDocIds(collectionIds).map(chunk -> {
        WriteBatch batch = helper.getFs().batch();
        chunk.stream().map(col -> helper.getBaseDocument().collection(col).document()).forEach(ref -> batch.set(ref, ImmutableMap.of("foo", "bar")));
        return batch.commit();
    }).collect(Collectors.toList())), FirestoreTestingHelper.flattenListList(), MoreExecutors.directExecutor()).get(10, TimeUnit.SECONDS);
    PCollection<String> actualCollectionIds = testPipeline.apply(Create.of("")).apply(getListCollectionIdsPTransform(testName.getMethodName())).apply(FirestoreIO.v1().read().listCollectionIds().withRpcQosOptions(rpcQosOptions).build());
    PAssert.that(actualCollectionIds).containsInAnyOrder(collectionIds);
    testPipeline.run(options);
}