Bug 8301: Fix some issues with testProducerRegistrations
[controller.git] / opendaylight / md-sal / sal-distributed-datastore / src / test / java / org / opendaylight / controller / cluster / datastore / IntegrationTestKit.java
index 9570e0e4aa7d5e36823e6ae3379b5cf6800097fe..1c4bfcfbb23dcf2d23faccdb6492d7f328656bf3 100644 (file)
@@ -58,10 +58,16 @@ public class IntegrationTestKit extends ShardTestKit {
 
     protected DatastoreContext.Builder datastoreContextBuilder;
     protected DatastoreSnapshot restoreFromSnapshot;
+    private final int commitTimeout;
 
     public IntegrationTestKit(final ActorSystem actorSystem, final Builder datastoreContextBuilder) {
+        this(actorSystem, datastoreContextBuilder, 7);
+    }
+
+    public IntegrationTestKit(final ActorSystem actorSystem, final Builder datastoreContextBuilder, int commitTimeout) {
         super(actorSystem);
         this.datastoreContextBuilder = datastoreContextBuilder;
+        this.commitTimeout = commitTimeout;
     }
 
     public DatastoreContext.Builder getDatastoreContextBuilder() {
@@ -122,18 +128,18 @@ public class IntegrationTestKit extends ShardTestKit {
         final ClusterWrapper cluster = new ClusterWrapperImpl(getSystem());
         final Configuration config = new ConfigurationImpl(moduleShardsConfig, modulesConfig);
 
-        datastoreContextBuilder.dataStoreName(typeName);
+        setDataStoreName(typeName);
 
         final DatastoreContext datastoreContext = datastoreContextBuilder.build();
         final DatastoreContextFactory mockContextFactory = Mockito.mock(DatastoreContextFactory.class);
         Mockito.doReturn(datastoreContext).when(mockContextFactory).getBaseDatastoreContext();
         Mockito.doReturn(datastoreContext).when(mockContextFactory).getShardDatastoreContext(Mockito.anyString());
 
-        final Constructor constructor = implementation.getDeclaredConstructor(
+        final Constructor<? extends AbstractDataStore> constructor = implementation.getDeclaredConstructor(
                 ActorSystem.class, ClusterWrapper.class, Configuration.class,
                 DatastoreContextFactory.class, DatastoreSnapshot.class);
 
-        final AbstractDataStore dataStore = (AbstractDataStore) constructor.newInstance(
+        final AbstractDataStore dataStore = constructor.newInstance(
                 getSystem(), cluster, config, mockContextFactory, restoreFromSnapshot);
 
         dataStore.onGlobalContextUpdated(schemaContext);
@@ -146,12 +152,22 @@ public class IntegrationTestKit extends ShardTestKit {
         return dataStore;
     }
 
+    private void setDataStoreName(String typeName) {
+        if ("config".equals(typeName)) {
+            datastoreContextBuilder.logicalStoreType(LogicalDatastoreType.CONFIGURATION);
+        } else if ("operational".equals(typeName)) {
+            datastoreContextBuilder.logicalStoreType(LogicalDatastoreType.OPERATIONAL);
+        } else {
+            datastoreContextBuilder.dataStoreName(typeName);
+        }
+    }
+
     public DistributedDataStore setupDistributedDataStoreWithoutConfig(final String typeName,
                                                                        final SchemaContext schemaContext) {
         final ClusterWrapper cluster = new ClusterWrapperImpl(getSystem());
         final ConfigurationImpl configuration = new ConfigurationImpl(new EmptyModuleShardConfigProvider());
 
-        getDatastoreContextBuilder().dataStoreName(typeName);
+        setDataStoreName(typeName);
 
         final DatastoreContext datastoreContext = getDatastoreContextBuilder().build();
 
@@ -174,7 +190,7 @@ public class IntegrationTestKit extends ShardTestKit {
         final ClusterWrapper cluster = new ClusterWrapperImpl(getSystem());
         final ConfigurationImpl configuration = new ConfigurationImpl(new EmptyModuleShardConfigProvider());
 
-        getDatastoreContextBuilder().dataStoreName(typeName);
+        setDataStoreName(typeName);
 
         final DatastoreContext datastoreContext =
                 getDatastoreContextBuilder().logicalStoreType(storeType).build();
@@ -334,7 +350,7 @@ public class IntegrationTestKit extends ShardTestKit {
     }
 
     public void doCommit(final DOMStoreThreePhaseCommitCohort cohort) throws Exception {
-        Boolean canCommit = cohort.canCommit().get(7, TimeUnit.SECONDS);
+        Boolean canCommit = cohort.canCommit().get(commitTimeout, TimeUnit.SECONDS);
         assertEquals("canCommit", true, canCommit);
         cohort.preCommit().get(5, TimeUnit.SECONDS);
         cohort.commit().get(5, TimeUnit.SECONDS);
@@ -342,7 +358,7 @@ public class IntegrationTestKit extends ShardTestKit {
 
     void doCommit(final ListenableFuture<Boolean> canCommitFuture, final DOMStoreThreePhaseCommitCohort cohort)
             throws Exception {
-        Boolean canCommit = canCommitFuture.get(7, TimeUnit.SECONDS);
+        Boolean canCommit = canCommitFuture.get(commitTimeout, TimeUnit.SECONDS);
         assertEquals("canCommit", true, canCommit);
         cohort.preCommit().get(5, TimeUnit.SECONDS);
         cohort.commit().get(5, TimeUnit.SECONDS);