X-Git-Url: https://git.opendaylight.org/gerrit/gitweb?a=blobdiff_plain;f=opendaylight%2Fmd-sal%2Fsal-distributed-datastore%2Fsrc%2Ftest%2Fjava%2Forg%2Fopendaylight%2Fcontroller%2Fcluster%2Fdatastore%2Futils%2FActorContextTest.java;h=093f0fa6df046ba333858117e385a2a971e12fe6;hb=refs%2Fchanges%2F47%2F29947%2F5;hp=efe73d38ad526e9a4c312949c79e0d5d64b7c7f8;hpb=ffc46de334c8a903844b9f4aff73dc68b2401659;p=controller.git

diff --git a/opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/utils/ActorContextTest.java b/opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/utils/ActorContextTest.java
index efe73d38ad..093f0fa6df 100644
--- a/opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/utils/ActorContextTest.java
+++ b/opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/utils/ActorContextTest.java
@@ -31,7 +31,6 @@ import akka.util.Timeout;
 import com.google.common.base.Optional;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
-import com.google.common.util.concurrent.Uninterruptibles;
 import com.typesafe.config.ConfigFactory;
 import java.util.Arrays;
 import java.util.Map;
@@ -43,6 +42,7 @@ import org.opendaylight.controller.cluster.datastore.AbstractActorTest;
 import org.opendaylight.controller.cluster.datastore.ClusterWrapper;
 import org.opendaylight.controller.cluster.datastore.DataStoreVersions;
 import org.opendaylight.controller.cluster.datastore.DatastoreContext;
+import org.opendaylight.controller.cluster.datastore.DatastoreContextFactory;
 import org.opendaylight.controller.cluster.datastore.config.Configuration;
 import org.opendaylight.controller.cluster.datastore.exceptions.NoShardLeaderException;
 import org.opendaylight.controller.cluster.datastore.exceptions.NotInitializedException;
@@ -54,6 +54,8 @@ import org.opendaylight.controller.cluster.datastore.messages.LocalShardFound;
 import org.opendaylight.controller.cluster.datastore.messages.LocalShardNotFound;
 import org.opendaylight.controller.cluster.datastore.messages.PrimaryShardInfo;
 import org.opendaylight.controller.cluster.datastore.messages.RemotePrimaryShardFound;
+import org.opendaylight.controller.cluster.raft.utils.EchoActor;
+import org.opendaylight.controller.cluster.raft.utils.MessageCollectorActor;
 import org.opendaylight.yangtools.yang.data.api.schema.tree.DataTree;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -374,9 +376,12 @@ public class ActorContextTest extends AbstractActorTest{
             DatastoreContext newContext = DatastoreContext.newBuilder().operationTimeoutInSeconds(6).
                     shardTransactionCommitTimeoutInSeconds(8).build();
 
-            actorContext.setDatastoreContext(newContext);
+            DatastoreContextFactory mockContextFactory = mock(DatastoreContextFactory.class);
+            Mockito.doReturn(newContext).when(mockContextFactory).getBaseDatastoreContext();
 
-            expectMsgClass(duration("5 seconds"), DatastoreContext.class);
+            actorContext.setDatastoreContext(mockContextFactory);
+
+            expectMsgClass(duration("5 seconds"), DatastoreContextFactory.class);
 
             Assert.assertSame("getDatastoreContext", newContext, actorContext.getDatastoreContext());
 
@@ -537,25 +542,9 @@ public class ActorContextTest extends AbstractActorTest{
 
             actorContext.broadcast(new TestMessage());
 
-            expectFirstMatching(shardActorRef1, TestMessage.class);
-            expectFirstMatching(shardActorRef2, TestMessage.class);
+            MessageCollectorActor.expectFirstMatching(shardActorRef1, TestMessage.class);
+            MessageCollectorActor.expectFirstMatching(shardActorRef2, TestMessage.class);
         }};
     }
 
-    private <T> T expectFirstMatching(ActorRef actor, Class<T> clazz) {
-        int count = 5000 / 50;
-        for(int i = 0; i < count; i++) {
-            try {
-                T message = (T) MessageCollectorActor.getFirstMatching(actor, clazz);
-                if(message != null) {
-                    return message;
-                }
-            } catch (Exception e) {}
-
-            Uninterruptibles.sleepUninterruptibly(50, TimeUnit.MILLISECONDS);
-        }
-
-        Assert.fail("Did not receive message of type " + clazz);
-        return null;
-    }
 }