Take snapshot after recovery on migrated messages

[controller.git] / opendaylight / md-sal / sal-akka-raft / src / test / java / org / opendaylight / controller / cluster / raft / MockRaftActor.java

diff --git a/opendaylight/md-sal/sal-akka-raft/src/test/java/org/opendaylight/controller/cluster/raft/MockRaftActor.java b/opendaylight/md-sal/sal-akka-raft/src/test/java/org/opendaylight/controller/cluster/raft/MockRaftActor.java
index 3ba664be7da313e3294fbc7c97988bf6c82018d6..88865d55bdf39a16891a3c59734aa4e1edf79ef2 100644 (file)
--- a/opendaylight/md-sal/sal-akka-raft/src/test/java/org/opendaylight/controller/cluster/raft/MockRaftActor.java
+++ b/opendaylight/md-sal/sal-akka-raft/src/test/java/org/opendaylight/controller/cluster/raft/MockRaftActor.java
@@ -52,7 +52,9 @@ public class MockRaftActor extends RaftActor implements RaftActorRecoveryCohort,
         state = new ArrayList<>();
         this.actorDelegate = mock(RaftActor.class);
         this.recoveryCohortDelegate = mock(RaftActorRecoveryCohort.class);
-        this.snapshotCohortDelegate = mock(RaftActorSnapshotCohort.class);
+
+        this.snapshotCohortDelegate = builder.snapshotCohort != null ? builder.snapshotCohort :
+            mock(RaftActorSnapshotCohort.class);
 
         if(builder.dataPersistenceProvider == null){
             setPersistence(builder.persistent.isPresent() ? builder.persistent.get() : true);
@@ -284,6 +286,7 @@ public class MockRaftActor extends RaftActor implements RaftActorRecoveryCohort,
         private Optional<Boolean> persistent = Optional.absent();
         private final Class<A> actorClass;
         private Function<Runnable, Void> pauseLeaderFunction;
+        private RaftActorSnapshotCohort snapshotCohort;
 
         protected AbstractBuilder(Class<A> actorClass) {
             this.actorClass = actorClass;
@@ -339,6 +342,11 @@ public class MockRaftActor extends RaftActor implements RaftActorRecoveryCohort,
             return self();
         }
 
+        public T snapshotCohort(RaftActorSnapshotCohort snapshotCohort) {
+            this.snapshotCohort = snapshotCohort;
+            return self();
+        }
+
         public Props props() {
             return Props.create(actorClass, this);
         }