Refactor ElectionTermImpl into separate class
[controller.git] / opendaylight / md-sal / sal-akka-raft / src / test / java / org / opendaylight / controller / cluster / raft / RaftActorTest.java
index 56bfc21f23c2047a19d5a49c49861ae1bf6eca5e..17a81ac3c39aa9c0a27743e1739892629a157ba8 100644 (file)
@@ -31,6 +31,7 @@ import akka.testkit.JavaTestKit;
 import akka.testkit.TestActorRef;
 import akka.util.Timeout;
 import com.google.common.base.Optional;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
 import com.google.common.util.concurrent.Uninterruptibles;
 import com.google.protobuf.ByteString;
@@ -53,24 +54,28 @@ import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.opendaylight.controller.cluster.DataPersistenceProvider;
+import org.opendaylight.controller.cluster.NonPersistentDataProvider;
 import org.opendaylight.controller.cluster.datastore.DataPersistenceProviderMonitor;
+import org.opendaylight.controller.cluster.notifications.LeaderStateChanged;
 import org.opendaylight.controller.cluster.notifications.RoleChanged;
+import org.opendaylight.controller.cluster.raft.RaftActor.UpdateElectionTerm;
+import org.opendaylight.controller.cluster.raft.base.messages.ApplyJournalEntries;
 import org.opendaylight.controller.cluster.raft.base.messages.ApplyLogEntries;
 import org.opendaylight.controller.cluster.raft.base.messages.ApplySnapshot;
 import org.opendaylight.controller.cluster.raft.base.messages.ApplyState;
-import org.opendaylight.controller.cluster.raft.base.messages.CaptureSnapshot;
 import org.opendaylight.controller.cluster.raft.base.messages.CaptureSnapshotReply;
 import org.opendaylight.controller.cluster.raft.base.messages.SendHeartBeat;
 import org.opendaylight.controller.cluster.raft.behaviors.Follower;
 import org.opendaylight.controller.cluster.raft.behaviors.Leader;
+import org.opendaylight.controller.cluster.raft.behaviors.RaftActorBehavior;
 import org.opendaylight.controller.cluster.raft.client.messages.FindLeader;
 import org.opendaylight.controller.cluster.raft.client.messages.FindLeaderReply;
 import org.opendaylight.controller.cluster.raft.messages.AppendEntries;
 import org.opendaylight.controller.cluster.raft.messages.AppendEntriesReply;
 import org.opendaylight.controller.cluster.raft.protobuff.client.messages.Payload;
+import org.opendaylight.controller.cluster.raft.utils.InMemoryJournal;
+import org.opendaylight.controller.cluster.raft.utils.InMemorySnapshotStore;
 import org.opendaylight.controller.cluster.raft.utils.MessageCollectorActor;
-import org.opendaylight.controller.cluster.raft.utils.MockAkkaJournal;
-import org.opendaylight.controller.cluster.raft.utils.MockSnapshotStore;
 import scala.concurrent.Await;
 import scala.concurrent.Future;
 import scala.concurrent.duration.Duration;
@@ -88,13 +93,12 @@ public class RaftActorTest extends AbstractActorTest {
     @After
     public void tearDown() throws Exception {
         factory.close();
-        MockAkkaJournal.clearJournal();
-        MockSnapshotStore.setMockSnapshot(null);
+        InMemoryJournal.clear();
+        InMemorySnapshotStore.clear();
     }
 
     public static class MockRaftActor extends RaftActor {
 
-        private final DataPersistenceProvider dataPersistenceProvider;
         private final RaftActor delegate;
         private final CountDownLatch recoveryComplete = new CountDownLatch(1);
         private final List<Object> state;
@@ -134,9 +138,9 @@ public class RaftActorTest extends AbstractActorTest {
             state = new ArrayList<>();
             this.delegate = mock(RaftActor.class);
             if(dataPersistenceProvider == null){
-                this.dataPersistenceProvider = new PersistentDataProvider();
+                setPersistence(true);
             } else {
-                this.dataPersistenceProvider = dataPersistenceProvider;
+                setPersistence(dataPersistenceProvider);
             }
         }
 
@@ -156,6 +160,16 @@ public class RaftActorTest extends AbstractActorTest {
             }
         }
 
+
+        public void waitUntilLeader(){
+            for(int i = 0;i < 10; i++){
+                if(isLeader()){
+                    break;
+                }
+                Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
+            }
+        }
+
         public List<Object> getState() {
             return state;
         }
@@ -175,9 +189,16 @@ public class RaftActorTest extends AbstractActorTest {
             return Props.create(new MockRaftActorCreator(peerAddresses, id, config, null, roleChangeNotifier));
         }
 
+        public static Props props(final String id, final Map<String, String> peerAddresses,
+                                  Optional<ConfigParams> config, ActorRef roleChangeNotifier,
+                                  DataPersistenceProvider dataPersistenceProvider){
+            return Props.create(new MockRaftActorCreator(peerAddresses, id, config, dataPersistenceProvider, roleChangeNotifier));
+        }
+
+
         @Override protected void applyState(ActorRef clientActor, String identifier, Object data) {
             delegate.applyState(clientActor, identifier, data);
-            LOG.info("applyState called");
+            LOG.info("{}: applyState called", persistenceId());
         }
 
         @Override
@@ -219,10 +240,12 @@ public class RaftActorTest extends AbstractActorTest {
         }
 
         @Override protected void createSnapshot() {
+            LOG.info("{}: createSnapshot called", persistenceId());
             delegate.createSnapshot();
         }
 
         @Override protected void applySnapshot(byte [] snapshot) {
+            LOG.info("{}: applySnapshot called", persistenceId());
             delegate.applySnapshot(snapshot);
         }
 
@@ -230,11 +253,6 @@ public class RaftActorTest extends AbstractActorTest {
             delegate.onStateChanged();
         }
 
-        @Override
-        protected DataPersistenceProvider persistence() {
-            return this.dataPersistenceProvider;
-        }
-
         @Override
         protected Optional<ActorRef> getRoleChangeNotifier() {
             return Optional.fromNullable(roleChangeNotifier);
@@ -270,7 +288,7 @@ public class RaftActorTest extends AbstractActorTest {
     }
 
 
-    private static class RaftActorTestKit extends JavaTestKit {
+    public static class RaftActorTestKit extends JavaTestKit {
         private final ActorRef raftActor;
 
         public RaftActorTestKit(ActorSystem actorSystem, String actorName) {
@@ -306,7 +324,7 @@ public class RaftActorTest extends AbstractActorTest {
             waitUntilLeader(raftActor);
         }
 
-        protected void waitUntilLeader(ActorRef actorRef) {
+        public static void waitUntilLeader(ActorRef actorRef) {
             FiniteDuration duration = Duration.create(100, TimeUnit.MILLISECONDS);
             for(int i = 0; i < 20 * 5; i++) {
                 Future<Object> future = Patterns.ask(actorRef, new FindLeader(), new Timeout(duration));
@@ -376,8 +394,7 @@ public class RaftActorTest extends AbstractActorTest {
             Snapshot snapshot = Snapshot.create(snapshotBytes.toByteArray(),
                     snapshotUnappliedEntries, lastIndexDuringSnapshotCapture, 1,
                     lastAppliedDuringSnapshotCapture, 1);
-            MockSnapshotStore.setMockSnapshot(snapshot);
-            MockSnapshotStore.setPersistenceId(persistenceId);
+            InMemorySnapshotStore.addSnapshot(persistenceId, snapshot);
 
             // add more entries after snapshot is taken
             List<ReplicatedLogEntry> entries = new ArrayList<>();
@@ -394,11 +411,11 @@ public class RaftActorTest extends AbstractActorTest {
             int lastAppliedToState = 5;
             int lastIndex = 7;
 
-            MockAkkaJournal.addToJournal(5, entry2);
+            InMemoryJournal.addEntry(persistenceId, 5, entry2);
             // 2 entries are applied to state besides the 4 entries in snapshot
-            MockAkkaJournal.addToJournal(6, new ApplyLogEntries(lastAppliedToState));
-            MockAkkaJournal.addToJournal(7, entry3);
-            MockAkkaJournal.addToJournal(8, entry4);
+            InMemoryJournal.addEntry(persistenceId, 6, new ApplyJournalEntries(lastAppliedToState));
+            InMemoryJournal.addEntry(persistenceId, 7, entry3);
+            InMemoryJournal.addEntry(persistenceId, 8, entry4);
 
             // kill the actor
             followerActor.tell(PoisonPill.getInstance(), null);
@@ -423,6 +440,46 @@ public class RaftActorTest extends AbstractActorTest {
         }};
     }
 
+    @Test
+    public void testRaftActorRecoveryWithPreLithuimApplyLogEntries() throws Exception {
+        new JavaTestKit(getSystem()) {{
+            String persistenceId = factory.generateActorId("leader-");
+
+            DefaultConfigParamsImpl config = new DefaultConfigParamsImpl();
+            config.setHeartBeatInterval(new FiniteDuration(1, TimeUnit.DAYS));
+
+            // Setup the persisted journal with some entries
+            ReplicatedLogEntry entry0 = new MockRaftActorContext.MockReplicatedLogEntry(1, 0,
+                    new MockRaftActorContext.MockPayload("zero"));
+            ReplicatedLogEntry entry1 = new MockRaftActorContext.MockReplicatedLogEntry(1, 1,
+                    new MockRaftActorContext.MockPayload("oen"));
+            ReplicatedLogEntry entry2 = new MockRaftActorContext.MockReplicatedLogEntry(1, 2,
+                    new MockRaftActorContext.MockPayload("two"));
+
+            long seqNr = 1;
+            InMemoryJournal.addEntry(persistenceId, seqNr++, entry0);
+            InMemoryJournal.addEntry(persistenceId, seqNr++, entry1);
+            InMemoryJournal.addEntry(persistenceId, seqNr++, new ApplyLogEntries(1));
+            InMemoryJournal.addEntry(persistenceId, seqNr++, entry2);
+
+            int lastAppliedToState = 1;
+            int lastIndex = 2;
+
+            //reinstate the actor
+            TestActorRef<MockRaftActor> leaderActor = factory.createTestActor(
+                    MockRaftActor.props(persistenceId, Collections.<String, String>emptyMap(),
+                            Optional.<ConfigParams>of(config)));
+
+            leaderActor.underlyingActor().waitForRecoveryComplete();
+
+            RaftActorContext context = leaderActor.underlyingActor().getRaftActorContext();
+            assertEquals("Journal log size", 3, context.getReplicatedLog().size());
+            assertEquals("Last index", lastIndex, context.getReplicatedLog().lastIndex());
+            assertEquals("Last applied", lastAppliedToState, context.getLastApplied());
+            assertEquals("Commit index", lastAppliedToState, context.getCommitIndex());
+        }};
+    }
+
     /**
      * This test verifies that when recovery is applicable (typically when persistence is true) the RaftActor does
      * process recovery messages
@@ -471,7 +528,7 @@ public class RaftActorTest extends AbstractActorTest {
 
                 assertEquals("add replicated log entry", 2, replicatedLog.size());
 
-                mockRaftActor.onReceiveRecover(new ApplyLogEntries(1));
+                mockRaftActor.onReceiveRecover(new ApplyJournalEntries(1));
 
                 assertEquals("commit index 1", 1, mockRaftActor.getRaftActorContext().getCommitIndex());
 
@@ -481,7 +538,7 @@ public class RaftActorTest extends AbstractActorTest {
 
                 assertEquals("remove log entries", 1, replicatedLog.size());
 
-                mockRaftActor.onReceiveRecover(new RaftActor.UpdateElectionTerm(10, "foobar"));
+                mockRaftActor.onReceiveRecover(new UpdateElectionTerm(10, "foobar"));
 
                 assertEquals("election term", 10, mockRaftActor.getRaftActorContext().getTermInformation().getCurrentTerm());
                 assertEquals("voted for", "foobar", mockRaftActor.getRaftActorContext().getTermInformation().getVotedFor());
@@ -538,7 +595,7 @@ public class RaftActorTest extends AbstractActorTest {
 
                 assertEquals("add replicated log entry", 0, replicatedLog.size());
 
-                mockRaftActor.onReceiveRecover(new ApplyLogEntries(1));
+                mockRaftActor.onReceiveRecover(new ApplyJournalEntries(1));
 
                 assertEquals("commit index -1", -1, mockRaftActor.getRaftActorContext().getCommitIndex());
 
@@ -546,7 +603,7 @@ public class RaftActorTest extends AbstractActorTest {
 
                 assertEquals("remove log entries", 0, replicatedLog.size());
 
-                mockRaftActor.onReceiveRecover(new RaftActor.UpdateElectionTerm(10, "foobar"));
+                mockRaftActor.onReceiveRecover(new UpdateElectionTerm(10, "foobar"));
 
                 assertNotEquals("election term", 10, mockRaftActor.getRaftActorContext().getTermInformation().getCurrentTerm());
                 assertNotEquals("voted for", "foobar", mockRaftActor.getRaftActorContext().getTermInformation().getVotedFor());
@@ -631,17 +688,19 @@ public class RaftActorTest extends AbstractActorTest {
 
                 mockRaftActor.waitForInitializeBehaviorComplete();
 
+                mockRaftActor.waitUntilLeader();
+
                 mockRaftActor.getReplicatedLog().appendAndPersist(new MockRaftActorContext.MockReplicatedLogEntry(1, 0, mock(Payload.class)));
 
                 mockRaftActor.getRaftActorContext().getReplicatedLog().removeFromAndPersist(0);
 
-                verify(dataPersistenceProvider, times(2)).persist(anyObject(), any(Procedure.class));
+                verify(dataPersistenceProvider, times(3)).persist(anyObject(), any(Procedure.class));
             }
         };
     }
 
     @Test
-    public void testApplyLogEntriesCallsDataPersistence() throws Exception {
+    public void testApplyJournalEntriesCallsDataPersistence() throws Exception {
         new JavaTestKit(getSystem()) {
             {
                 String persistenceId = factory.generateActorId("leader-");
@@ -659,9 +718,11 @@ public class RaftActorTest extends AbstractActorTest {
 
                 mockRaftActor.waitForInitializeBehaviorComplete();
 
-                mockRaftActor.onReceiveCommand(new ApplyLogEntries(10));
+                mockRaftActor.waitUntilLeader();
 
-                verify(dataPersistenceProvider, times(1)).persist(anyObject(), any(Procedure.class));
+                mockRaftActor.onReceiveCommand(new ApplyJournalEntries(10));
+
+                verify(dataPersistenceProvider, times(2)).persist(anyObject(), any(Procedure.class));
 
             }
 
@@ -694,10 +755,12 @@ public class RaftActorTest extends AbstractActorTest {
                         new MockRaftActorContext.MockPayload("C"),
                         new MockRaftActorContext.MockPayload("D")));
 
-                mockRaftActor.onReceiveCommand(new CaptureSnapshot(-1, 1,-1, 1, -1, 1));
-
                 RaftActorContext raftActorContext = mockRaftActor.getRaftActorContext();
 
+                raftActorContext.getSnapshotManager().capture(
+                        new MockRaftActorContext.MockReplicatedLogEntry(1, -1,
+                                new MockRaftActorContext.MockPayload("D")), -1);
+
                 mockRaftActor.setCurrentBehavior(new Leader(raftActorContext));
 
                 mockRaftActor.onReceiveCommand(new CaptureSnapshotReply(snapshotBytes.toByteArray()));
@@ -721,17 +784,18 @@ public class RaftActorTest extends AbstractActorTest {
                 DataPersistenceProvider dataPersistenceProvider = mock(DataPersistenceProvider.class);
 
                 TestActorRef<MockRaftActor> mockActorRef = factory.createTestActor(MockRaftActor.props(persistenceId,
-                        Collections.<String, String>emptyMap(), Optional.<ConfigParams>of(config), dataPersistenceProvider), persistenceId);
+                        ImmutableMap.of("leader", "fake/path"), Optional.<ConfigParams>of(config), dataPersistenceProvider), persistenceId);
 
                 MockRaftActor mockRaftActor = mockActorRef.underlyingActor();
 
                 mockRaftActor.waitForInitializeBehaviorComplete();
+                MockRaftActorContext.MockReplicatedLogEntry lastEntry = new MockRaftActorContext.MockReplicatedLogEntry(1, 4, mock(Payload.class));
 
                 mockRaftActor.getReplicatedLog().append(new MockRaftActorContext.MockReplicatedLogEntry(1, 0, mock(Payload.class)));
                 mockRaftActor.getReplicatedLog().append(new MockRaftActorContext.MockReplicatedLogEntry(1, 1, mock(Payload.class)));
                 mockRaftActor.getReplicatedLog().append(new MockRaftActorContext.MockReplicatedLogEntry(1, 2, mock(Payload.class)));
                 mockRaftActor.getReplicatedLog().append(new MockRaftActorContext.MockReplicatedLogEntry(1, 3, mock(Payload.class)));
-                mockRaftActor.getReplicatedLog().append(new MockRaftActorContext.MockReplicatedLogEntry(1, 4, mock(Payload.class)));
+                mockRaftActor.getReplicatedLog().append(lastEntry);
 
                 ByteString snapshotBytes = fromObject(Arrays.asList(
                         new MockRaftActorContext.MockPayload("A"),
@@ -743,7 +807,8 @@ public class RaftActorTest extends AbstractActorTest {
                 mockRaftActor.setCurrentBehavior(new Follower(raftActorContext));
 
                 long replicatedToAllIndex = 1;
-                mockRaftActor.onReceiveCommand(new CaptureSnapshot(-1, 1, 2, 1, replicatedToAllIndex, 1));
+
+                mockRaftActor.getRaftActorContext().getSnapshotManager().capture(lastEntry, replicatedToAllIndex);
 
                 verify(mockRaftActor.delegate).createSnapshot();
 
@@ -885,7 +950,9 @@ public class RaftActorTest extends AbstractActorTest {
 
                 mockRaftActor.setCurrentBehavior(new Leader(raftActorContext));
 
-                mockRaftActor.onReceiveCommand(new CaptureSnapshot(-1, 1, -1, 1, -1, 1));
+                raftActorContext.getSnapshotManager().capture(
+                        new MockRaftActorContext.MockReplicatedLogEntry(1, 1,
+                                new MockRaftActorContext.MockPayload("D")), 1);
 
                 mockRaftActor.onReceiveCommand(new CaptureSnapshotReply(snapshotBytes.toByteArray()));
 
@@ -900,7 +967,83 @@ public class RaftActorTest extends AbstractActorTest {
     }
 
     @Test
-    public void testRaftRoleChangeNotifier() throws Exception {
+    public void testRaftRoleChangeNotifierWhenRaftActorHasNoPeers() throws Exception {
+        new JavaTestKit(getSystem()) {{
+            TestActorRef<MessageCollectorActor> notifierActor = factory.createTestActor(
+                    Props.create(MessageCollectorActor.class));
+            MessageCollectorActor.waitUntilReady(notifierActor);
+
+            DefaultConfigParamsImpl config = new DefaultConfigParamsImpl();
+            long heartBeatInterval = 100;
+            config.setHeartBeatInterval(FiniteDuration.create(heartBeatInterval, TimeUnit.MILLISECONDS));
+            config.setElectionTimeoutFactor(20);
+
+            String persistenceId = factory.generateActorId("notifier-");
+
+            TestActorRef<MockRaftActor> raftActorRef = factory.createTestActor(MockRaftActor.props(persistenceId,
+                    Collections.<String, String>emptyMap(), Optional.<ConfigParams>of(config), notifierActor,
+                    new NonPersistentDataProvider()), persistenceId);
+
+            List<RoleChanged> matches =  MessageCollectorActor.expectMatching(notifierActor, RoleChanged.class, 3);
+
+
+            // check if the notifier got a role change from null to Follower
+            RoleChanged raftRoleChanged = matches.get(0);
+            assertEquals(persistenceId, raftRoleChanged.getMemberId());
+            assertNull(raftRoleChanged.getOldRole());
+            assertEquals(RaftState.Follower.name(), raftRoleChanged.getNewRole());
+
+            // check if the notifier got a role change from Follower to Candidate
+            raftRoleChanged = matches.get(1);
+            assertEquals(persistenceId, raftRoleChanged.getMemberId());
+            assertEquals(RaftState.Follower.name(), raftRoleChanged.getOldRole());
+            assertEquals(RaftState.Candidate.name(), raftRoleChanged.getNewRole());
+
+            // check if the notifier got a role change from Candidate to Leader
+            raftRoleChanged = matches.get(2);
+            assertEquals(persistenceId, raftRoleChanged.getMemberId());
+            assertEquals(RaftState.Candidate.name(), raftRoleChanged.getOldRole());
+            assertEquals(RaftState.Leader.name(), raftRoleChanged.getNewRole());
+
+            LeaderStateChanged leaderStateChange = MessageCollectorActor.expectFirstMatching(
+                    notifierActor, LeaderStateChanged.class);
+
+            assertEquals(raftRoleChanged.getMemberId(), leaderStateChange.getLeaderId());
+
+            notifierActor.underlyingActor().clear();
+
+            MockRaftActor raftActor = raftActorRef.underlyingActor();
+            final String newLeaderId = "new-leader";
+            Follower follower = new Follower(raftActor.getRaftActorContext()) {
+                @Override
+                public RaftActorBehavior handleMessage(ActorRef sender, Object message) {
+                    leaderId = newLeaderId;
+                    return this;
+                }
+            };
+
+            raftActor.changeCurrentBehavior(follower);
+
+            leaderStateChange = MessageCollectorActor.expectFirstMatching(notifierActor, LeaderStateChanged.class);
+            assertEquals(persistenceId, leaderStateChange.getMemberId());
+            assertEquals(null, leaderStateChange.getLeaderId());
+
+            raftRoleChanged = MessageCollectorActor.expectFirstMatching(notifierActor, RoleChanged.class);
+            assertEquals(RaftState.Leader.name(), raftRoleChanged.getOldRole());
+            assertEquals(RaftState.Follower.name(), raftRoleChanged.getNewRole());
+
+            notifierActor.underlyingActor().clear();
+
+            raftActor.handleCommand("any");
+
+            leaderStateChange = MessageCollectorActor.expectFirstMatching(notifierActor, LeaderStateChanged.class);
+            assertEquals(persistenceId, leaderStateChange.getMemberId());
+            assertEquals(newLeaderId, leaderStateChange.getLeaderId());
+        }};
+    }
+
+    @Test
+    public void testRaftRoleChangeNotifierWhenRaftActorHasPeers() throws Exception {
         new JavaTestKit(getSystem()) {{
             ActorRef notifierActor = factory.createActor(Props.create(MessageCollectorActor.class));
             MessageCollectorActor.waitUntilReady(notifierActor);
@@ -912,8 +1055,8 @@ public class RaftActorTest extends AbstractActorTest {
 
             String persistenceId = factory.generateActorId("notifier-");
 
-            factory.createTestActor(MockRaftActor.props(persistenceId,
-                    Collections.<String, String>emptyMap(), Optional.<ConfigParams>of(config), notifierActor), persistenceId);
+            factory.createActor(MockRaftActor.props(persistenceId,
+                    ImmutableMap.of("leader", "fake/path"), Optional.<ConfigParams>of(config), notifierActor), persistenceId);
 
             List<RoleChanged> matches =  null;
             for(int i = 0; i < 5000 / heartBeatInterval; i++) {
@@ -925,7 +1068,7 @@ public class RaftActorTest extends AbstractActorTest {
                 Uninterruptibles.sleepUninterruptibly(heartBeatInterval, TimeUnit.MILLISECONDS);
             }
 
-            assertEquals(3, matches.size());
+            assertEquals(2, matches.size());
 
             // check if the notifier got a role change from null to Follower
             RoleChanged raftRoleChanged = matches.get(0);
@@ -939,11 +1082,6 @@ public class RaftActorTest extends AbstractActorTest {
             assertEquals(RaftState.Follower.name(), raftRoleChanged.getOldRole());
             assertEquals(RaftState.Candidate.name(), raftRoleChanged.getNewRole());
 
-            // check if the notifier got a role change from Candidate to Leader
-            raftRoleChanged = matches.get(2);
-            assertEquals(persistenceId, raftRoleChanged.getMemberId());
-            assertEquals(RaftState.Candidate.name(), raftRoleChanged.getOldRole());
-            assertEquals(RaftState.Leader.name(), raftRoleChanged.getNewRole());
         }};
     }
 
@@ -989,9 +1127,10 @@ public class RaftActorTest extends AbstractActorTest {
 
                 assertEquals(8, leaderActor.getReplicatedLog().size());
 
-                leaderActor.onReceiveCommand(new CaptureSnapshot(6, 1, 4, 1, 4, 1));
+                leaderActor.getRaftActorContext().getSnapshotManager()
+                        .capture(new MockRaftActorContext.MockReplicatedLogEntry(1, 6,
+                                new MockRaftActorContext.MockPayload("x")), 4);
 
-                leaderActor.getRaftActorContext().setSnapshotCaptureInitiated(true);
                 verify(leaderActor.delegate).createSnapshot();
 
                 assertEquals(8, leaderActor.getReplicatedLog().size());
@@ -1017,8 +1156,14 @@ public class RaftActorTest extends AbstractActorTest {
                         new MockRaftActorContext.MockPayload("foo-2"),
                         new MockRaftActorContext.MockPayload("foo-3"),
                         new MockRaftActorContext.MockPayload("foo-4")));
-                leaderActor.onReceiveCommand(new CaptureSnapshotReply(snapshotBytes.toByteArray()));
-                assertFalse(leaderActor.getRaftActorContext().isSnapshotCaptureInitiated());
+
+                leaderActor.getRaftActorContext().getSnapshotManager().persist(new NonPersistentDataProvider()
+                        , snapshotBytes.toByteArray(), leader, Runtime.getRuntime().totalMemory());
+
+                assertFalse(leaderActor.getRaftActorContext().getSnapshotManager().isCapturing());
+
+                // The commit is needed to complete the snapshot creation process
+                leaderActor.getRaftActorContext().getSnapshotManager().commit(new NonPersistentDataProvider(), -1);
 
                 // capture snapshot reply should remove the snapshotted entries only
                 assertEquals(3, leaderActor.getReplicatedLog().size());
@@ -1081,9 +1226,10 @@ public class RaftActorTest extends AbstractActorTest {
                 assertEquals(6, followerActor.getReplicatedLog().size());
 
                 //snapshot on 4
-                followerActor.onReceiveCommand(new CaptureSnapshot(5, 1, 4, 1, 4, 1));
+                followerActor.getRaftActorContext().getSnapshotManager().capture(
+                        new MockRaftActorContext.MockReplicatedLogEntry(1, 5,
+                                new MockRaftActorContext.MockPayload("D")), 4);
 
-                followerActor.getRaftActorContext().setSnapshotCaptureInitiated(true);
                 verify(followerActor.delegate).createSnapshot();
 
                 assertEquals(6, followerActor.getReplicatedLog().size());
@@ -1118,7 +1264,10 @@ public class RaftActorTest extends AbstractActorTest {
                         new MockRaftActorContext.MockPayload("foo-3"),
                         new MockRaftActorContext.MockPayload("foo-4")));
                 followerActor.onReceiveCommand(new CaptureSnapshotReply(snapshotBytes.toByteArray()));
-                assertFalse(followerActor.getRaftActorContext().isSnapshotCaptureInitiated());
+                assertFalse(followerActor.getRaftActorContext().getSnapshotManager().isCapturing());
+
+                // The commit is needed to complete the snapshot creation process
+                followerActor.getRaftActorContext().getSnapshotManager().commit(new NonPersistentDataProvider(), -1);
 
                 // capture snapshot reply should remove the snapshotted entries only till replicatedToAllIndex
                 assertEquals(3, followerActor.getReplicatedLog().size()); //indexes 5,6,7 left in the log
@@ -1216,7 +1365,7 @@ public class RaftActorTest extends AbstractActorTest {
                         new MockRaftActorContext.MockPayload("foo-3"),
                         new MockRaftActorContext.MockPayload("foo-4")));
                 leaderActor.onReceiveCommand(new CaptureSnapshotReply(snapshotBytes.toByteArray()));
-                assertFalse(leaderActor.getRaftActorContext().isSnapshotCaptureInitiated());
+                assertFalse(leaderActor.getRaftActorContext().getSnapshotManager().isCapturing());
 
                 assertEquals("Real snapshot didn't clear the log till replicatedToAllIndex", 0, leaderActor.getReplicatedLog().size());
 
@@ -1227,38 +1376,6 @@ public class RaftActorTest extends AbstractActorTest {
         };
     }
 
-
-    private static class NonPersistentProvider implements DataPersistenceProvider {
-        @Override
-        public boolean isRecoveryApplicable() {
-            return false;
-        }
-
-        @Override
-        public <T> void persist(T o, Procedure<T> procedure) {
-            try {
-                procedure.apply(o);
-            } catch (Exception e) {
-                e.printStackTrace();
-            }
-        }
-
-        @Override
-        public void saveSnapshot(Object o) {
-
-        }
-
-        @Override
-        public void deleteSnapshots(SnapshotSelectionCriteria criteria) {
-
-        }
-
-        @Override
-        public void deleteMessages(long sequenceNumber) {
-
-        }
-    }
-
     @Test
     public void testRealSnapshotWhenReplicatedToAllIndexMinusOne() throws Exception {
         new JavaTestKit(getSystem()) {{
@@ -1268,7 +1385,7 @@ public class RaftActorTest extends AbstractActorTest {
             config.setIsolatedLeaderCheckInterval(new FiniteDuration(1, TimeUnit.DAYS));
             config.setSnapshotBatchCount(5);
 
-            DataPersistenceProvider dataPersistenceProvider = new NonPersistentProvider();
+            DataPersistenceProvider dataPersistenceProvider = new NonPersistentDataProvider();
 
             Map<String, String> peerAddresses = new HashMap<>();
 
@@ -1300,7 +1417,7 @@ public class RaftActorTest extends AbstractActorTest {
 
             // Trimming log in this scenario is a no-op
             assertEquals(-1, leaderActor.getReplicatedLog().getSnapshotIndex());
-            assertFalse(leaderActor.getRaftActorContext().isSnapshotCaptureInitiated());
+            assertFalse(leaderActor.getRaftActorContext().getSnapshotManager().isCapturing());
             assertEquals(-1, leader.getReplicatedToAllIndex());
 
         }};
@@ -1315,7 +1432,7 @@ public class RaftActorTest extends AbstractActorTest {
             config.setIsolatedLeaderCheckInterval(new FiniteDuration(1, TimeUnit.DAYS));
             config.setSnapshotBatchCount(5);
 
-            DataPersistenceProvider dataPersistenceProvider = new NonPersistentProvider();
+            DataPersistenceProvider dataPersistenceProvider = new NonPersistentDataProvider();
 
             Map<String, String> peerAddresses = new HashMap<>();
 
@@ -1343,7 +1460,7 @@ public class RaftActorTest extends AbstractActorTest {
 
             // Trimming log in this scenario is a no-op
             assertEquals(3, leaderActor.getReplicatedLog().getSnapshotIndex());
-            assertFalse(leaderActor.getRaftActorContext().isSnapshotCaptureInitiated());
+            assertFalse(leaderActor.getRaftActorContext().getSnapshotManager().isCapturing());
             assertEquals(3, leader.getReplicatedToAllIndex());
 
         }};