FeatureRowToStoreAllocator extends PTransform<PCollection<FeatureRowProto.FeatureRow>, PCollectionTuple> { public static Builder newBuilder() { return new AutoValue_FeatureRowToStoreAllocator.Builder(); }  abstract List<StoreProto.Store> getStores(); abstract Map<StoreProto.Store, TupleTag<FeatureRowProto.FeatureRow>> getStoreTags(); static Builder newBuilder(); @Override PCollectionTuple expand(PCollection<FeatureRowProto.FeatureRow> input);  }
@Test public void featureRowShouldBeAllocatedToStoreTagsAccordingToSubscription() { StoreProto.Store bqOnlyStore = newStore("bq*"); StoreProto.Store redisOnlyStore = newStore("redis*"); StoreProto.Store anyStore = newStore("*"); Map<StoreProto.Store, TupleTag<FeatureRowProto.FeatureRow>> storeTags = ImmutableMap.of( bqOnlyStore, new TupleTag<>(), redisOnlyStore, new TupleTag<>(), anyStore, new TupleTag<>()); PCollectionTuple allocatedRows = p.apply( Create.of( FeatureRowProto.FeatureRow.newBuilder().setFeatureSet("project/bq_1").build(), FeatureRowProto.FeatureRow.newBuilder().setFeatureSet("project/bq_2").build(), FeatureRowProto.FeatureRow.newBuilder() .setFeatureSet("project/redis_1") .build(), FeatureRowProto.FeatureRow.newBuilder() .setFeatureSet("project/redis_2") .build(), FeatureRowProto.FeatureRow.newBuilder() .setFeatureSet("project/redis_3") .build())) .apply( FeatureRowToStoreAllocator.newBuilder() .setStoreTags(storeTags) .setStores(ImmutableList.of(bqOnlyStore, redisOnlyStore, anyStore)) .build()); PAssert.that( allocatedRows .get(storeTags.get(bqOnlyStore)) .setCoder(ProtoCoder.of(FeatureRowProto.FeatureRow.class)) .apply("CountBq", Count.globally())) .containsInAnyOrder(2L); PAssert.that( allocatedRows .get(storeTags.get(redisOnlyStore)) .setCoder(ProtoCoder.of(FeatureRowProto.FeatureRow.class)) .apply("CountRedis", Count.globally())) .containsInAnyOrder(3L); PAssert.that( allocatedRows .get(storeTags.get(anyStore)) .setCoder(ProtoCoder.of(FeatureRowProto.FeatureRow.class)) .apply("CountAll", Count.globally())) .containsInAnyOrder(5L); p.run(); }
@Test public void featureRowShouldBeAllocatedToStoreTagsAccordingToSubscriptionBlacklist() { Subscription subscription1 = Subscription.newBuilder().setProject("*").setName("*").build(); Subscription subscription2 = Subscription.newBuilder().setProject("project1").setName("fs_2").build(); Subscription subscription3 = Subscription.newBuilder().setProject("project1").setName("fs_1").setExclude(true).build(); Subscription subscription4 = Subscription.newBuilder().setProject("project2").setName("*").setExclude(true).build(); List<Subscription> testStoreSubscriptions1 = Arrays.asList(subscription1, subscription2, subscription3); StoreProto.Store testStore1 = newStore(testStoreSubscriptions1); List<Subscription> testStoreSubscriptions2 = Arrays.asList(subscription1, subscription4); StoreProto.Store testStore2 = newStore(testStoreSubscriptions2); Map<StoreProto.Store, TupleTag<FeatureRowProto.FeatureRow>> storeTags = ImmutableMap.of( testStore1, new TupleTag<>(), testStore2, new TupleTag<>()); PCollectionTuple allocatedRows = p.apply( Create.of( FeatureRowProto.FeatureRow.newBuilder().setFeatureSet("project1/fs_1").build(), FeatureRowProto.FeatureRow.newBuilder().setFeatureSet("project2/fs_1").build(), FeatureRowProto.FeatureRow.newBuilder().setFeatureSet("project2/fs_2").build())) .apply( FeatureRowToStoreAllocator.newBuilder() .setStoreTags(storeTags) .setStores(ImmutableList.of(testStore1, testStore2)) .build()); PAssert.that( allocatedRows .get(storeTags.get(testStore1)) .setCoder(ProtoCoder.of(FeatureRowProto.FeatureRow.class)) .apply("CountStore1", Count.globally())) .containsInAnyOrder(2L); PAssert.that( allocatedRows .get(storeTags.get(testStore2)) .setCoder(ProtoCoder.of(FeatureRowProto.FeatureRow.class)) .apply("CountStore2", Count.globally())) .containsInAnyOrder(1L); p.run(); }