Examples with SplitEntities google.registry.beam.datastore.BulkDeleteDatastorePipeline.SplitEntities used on opensource projects

Example 1 with SplitEntities

use of google.registry.beam.datastore.BulkDeleteDatastorePipeline.SplitEntities in project nomulus by google.

the class BulkDeleteDatastorePipelineTest method splitEntitiesByKind.

@Test
void splitEntitiesByKind() {
    TupleTagList tags = getDeletionTags(2);
    PCollection<String> kinds = testPipeline.apply("InjectKinds", Create.of("A", "B"));
    PCollectionView<Map<String, TupleTag<Entity>>> kindToTagMapping = BulkDeleteDatastorePipeline.mapKindsToDeletionTags(kinds, tags).apply(View.asMap());
    Entity entityA = createTestEntity("A", 1);
    Entity entityB = createTestEntity("B", 2);
    PCollection<Entity> entities = testPipeline.apply("InjectEntities", Create.of(entityA, entityB));
    PCollectionTuple allCollections = entities.apply("SplitByKind", ParDo.of(new SplitEntities(kindToTagMapping)).withSideInputs(kindToTagMapping).withOutputTags(getOneDeletionTag("placeholder"), tags));
    PAssert.that(allCollections.get((TupleTag<Entity>) tags.get(0))).containsInAnyOrder(entityA);
    PAssert.that(allCollections.get((TupleTag<Entity>) tags.get(1))).containsInAnyOrder(entityB);
    testPipeline.run();
}