Merge "Remove isCloseMsg check for each rpc"
[controller.git] / opendaylight / md-sal / sal-akka-raft / src / test / java / org / opendaylight / controller / cluster / raft / MockRaftActor.java
index 3275e8d6c480e20d8ef5ee06914c211c03679049..53110b3583a4f7353e3d17b1ec7ac91bf8ec7891 100644 (file)
@@ -37,6 +37,7 @@ public class MockRaftActor extends RaftActor implements RaftActorRecoveryCohort,
     private ActorRef roleChangeNotifier;
     private final CountDownLatch initializeBehaviorComplete = new CountDownLatch(1);
     private RaftActorRecoverySupport raftActorRecoverySupport;
+    private RaftActorSnapshotMessageSupport snapshotMessageSupport;
 
     public static final class MockRaftActorCreator implements Creator<MockRaftActor> {
         private static final long serialVersionUID = 1L;
@@ -45,6 +46,7 @@ public class MockRaftActor extends RaftActor implements RaftActorRecoveryCohort,
         private final Optional<ConfigParams> config;
         private final DataPersistenceProvider dataPersistenceProvider;
         private final ActorRef roleChangeNotifier;
+        private RaftActorSnapshotMessageSupport snapshotMessageSupport;
 
         private MockRaftActorCreator(Map<String, String> peerAddresses, String id,
             Optional<ConfigParams> config, DataPersistenceProvider dataPersistenceProvider,
@@ -61,6 +63,7 @@ public class MockRaftActor extends RaftActor implements RaftActorRecoveryCohort,
             MockRaftActor mockRaftActor = new MockRaftActor(id, peerAddresses, config,
                 dataPersistenceProvider);
             mockRaftActor.roleChangeNotifier = this.roleChangeNotifier;
+            mockRaftActor.snapshotMessageSupport = snapshotMessageSupport;
             return mockRaftActor;
         }
     }
@@ -88,6 +91,11 @@ public class MockRaftActor extends RaftActor implements RaftActorRecoveryCohort,
         return raftActorRecoverySupport != null ? raftActorRecoverySupport : super.newRaftActorRecoverySupport();
     }
 
+    @Override
+    protected RaftActorSnapshotMessageSupport newRaftActorSnapshotMessageSupport() {
+        return snapshotMessageSupport != null ? snapshotMessageSupport : super.newRaftActorSnapshotMessageSupport();
+    }
+
     public void waitForRecoveryComplete() {
         try {
             assertEquals("Recovery complete", true, recoveryComplete.await(5,  TimeUnit.SECONDS));
@@ -123,6 +131,13 @@ public class MockRaftActor extends RaftActor implements RaftActorRecoveryCohort,
         return Props.create(new MockRaftActorCreator(peerAddresses, id, config, null, null));
     }
 
+    public static Props props(final String id, final Map<String, String> peerAddresses,
+            Optional<ConfigParams> config, RaftActorSnapshotMessageSupport snapshotMessageSupport){
+        MockRaftActorCreator creator = new MockRaftActorCreator(peerAddresses, id, config, null, null);
+        creator.snapshotMessageSupport = snapshotMessageSupport;
+        return Props.create(creator);
+    }
+
     public static Props props(final String id, final Map<String, String> peerAddresses,
                               Optional<ConfigParams> config, DataPersistenceProvider dataPersistenceProvider){
         return Props.create(new MockRaftActorCreator(peerAddresses, id, config, dataPersistenceProvider, null));