Examples with ForwardingDatasetFramework co.cask.cdap.data2.dataset2.ForwardingDatasetFramework used on opensource projects

Example 1 with ForwardingDatasetFramework

use of co.cask.cdap.data2.dataset2.ForwardingDatasetFramework in project cdap by caskdata.

the class UsageRegistryTest method testUsageRegistry.

@Test
public void testUsageRegistry() {
    // instantiate a usage registry
    UsageRegistry registry = new DefaultUsageRegistry(new TransactionExecutorFactory() {

        @Override
        public TransactionExecutor createExecutor(Iterable<TransactionAware> iterable) {
            return dsFrameworkUtil.newInMemoryTransactionExecutor(iterable);
        }
    }, new ForwardingDatasetFramework(dsFrameworkUtil.getFramework()) {

        @Nullable
        @Override
        public <T extends Dataset> T getDataset(DatasetId datasetInstanceId, Map<String, String> arguments, @Nullable ClassLoader classLoader) throws DatasetManagementException, IOException {
            T t = super.getDataset(datasetInstanceId, arguments, classLoader);
            if (t instanceof UsageDataset) {
                @SuppressWarnings("unchecked") T t1 = (T) new WrappedUsageDataset((UsageDataset) t);
                return t1;
            }
            return t;
        }
    });
    // register usage for a stream and a dataset for single and multiple "owners", including a non-program
    registry.register(flow11, datasetInstance1);
    registry.register(flow12, stream1);
    registry.registerAll(ImmutableList.of(flow21, flow22), datasetInstance2);
    registry.registerAll(ImmutableList.of(flow21, flow22), stream1);
    int count = WrappedUsageDataset.registerCount;
    // validate usage
    Assert.assertEquals(ImmutableSet.of(datasetInstance1), registry.getDatasets(flow11));
    Assert.assertEquals(ImmutableSet.of(stream1), registry.getStreams(flow12));
    Assert.assertEquals(ImmutableSet.of(datasetInstance2), registry.getDatasets(flow21));
    Assert.assertEquals(ImmutableSet.of(datasetInstance2), registry.getDatasets(flow22));
    Assert.assertEquals(ImmutableSet.of(stream1), registry.getStreams(flow21));
    Assert.assertEquals(ImmutableSet.of(stream1), registry.getStreams(flow22));
    Assert.assertEquals(ImmutableSet.of(flow11), registry.getPrograms(datasetInstance1));
    Assert.assertEquals(ImmutableSet.of(flow21, flow22), registry.getPrograms(datasetInstance2));
    Assert.assertEquals(ImmutableSet.of(flow12, flow21, flow22), registry.getPrograms(stream1));
    // register datasets again
    registry.register(flow11, datasetInstance1);
    registry.registerAll(ImmutableList.of(flow21, flow22), datasetInstance2);
    // validate that this does re-register previous usages (DefaultUsageRegistry no longer avoids re-registration)
    count += 3;
    Assert.assertEquals(count, WrappedUsageDataset.registerCount);
    // validate usage
    Assert.assertEquals(ImmutableSet.of(datasetInstance1), registry.getDatasets(flow11));
    Assert.assertEquals(ImmutableSet.of(stream1), registry.getStreams(flow12));
    Assert.assertEquals(ImmutableSet.of(datasetInstance2), registry.getDatasets(flow21));
    Assert.assertEquals(ImmutableSet.of(datasetInstance2), registry.getDatasets(flow22));
    Assert.assertEquals(ImmutableSet.of(stream1), registry.getStreams(flow21));
    Assert.assertEquals(ImmutableSet.of(stream1), registry.getStreams(flow22));
    Assert.assertEquals(ImmutableSet.of(flow11), registry.getPrograms(datasetInstance1));
    Assert.assertEquals(ImmutableSet.of(flow21, flow22), registry.getPrograms(datasetInstance2));
    Assert.assertEquals(ImmutableSet.of(flow12, flow21, flow22), registry.getPrograms(stream1));
    // unregister app
    registry.unregister(flow11.getParent());
    // validate usage for that app is gone
    Assert.assertEquals(ImmutableSet.of(), registry.getDatasets(flow11));
    Assert.assertEquals(ImmutableSet.of(), registry.getStreams(flow12));
    Assert.assertEquals(ImmutableSet.of(datasetInstance2), registry.getDatasets(flow21));
    Assert.assertEquals(ImmutableSet.of(datasetInstance2), registry.getDatasets(flow22));
    Assert.assertEquals(ImmutableSet.of(stream1), registry.getStreams(flow21));
    Assert.assertEquals(ImmutableSet.of(stream1), registry.getStreams(flow22));
    Assert.assertEquals(ImmutableSet.of(), registry.getPrograms(datasetInstance1));
    Assert.assertEquals(ImmutableSet.of(flow21, flow22), registry.getPrograms(datasetInstance2));
    Assert.assertEquals(ImmutableSet.of(flow21, flow22), registry.getPrograms(stream1));
    // register application 1 again
    registry.register(flow11, datasetInstance1);
    registry.register(flow12, stream1);
    // validate it was re-registered
    Assert.assertEquals(ImmutableSet.of(datasetInstance1), registry.getDatasets(flow11));
    Assert.assertEquals(ImmutableSet.of(stream1), registry.getStreams(flow12));
    Assert.assertEquals(ImmutableSet.of(datasetInstance2), registry.getDatasets(flow21));
    Assert.assertEquals(ImmutableSet.of(datasetInstance2), registry.getDatasets(flow22));
    Assert.assertEquals(ImmutableSet.of(stream1), registry.getStreams(flow21));
    Assert.assertEquals(ImmutableSet.of(stream1), registry.getStreams(flow22));
    Assert.assertEquals(ImmutableSet.of(flow11), registry.getPrograms(datasetInstance1));
    Assert.assertEquals(ImmutableSet.of(flow21, flow22), registry.getPrograms(datasetInstance2));
    Assert.assertEquals(ImmutableSet.of(flow12, flow21, flow22), registry.getPrograms(stream1));
    // validate that this actually re-registered previous usages (through code in wrapped usage dataset)
    Assert.assertEquals(count + 2, WrappedUsageDataset.registerCount);
}