Merge "Removing { } from NormalizedNodeJsonBodyWriter"
[controller.git] / opendaylight / md-sal / sal-distributed-datastore / src / test / java / org / opendaylight / controller / cluster / datastore / ShardTest.java
index 926cef6ba53ea2805f4310e2d9a55b59371f4094..4b0651a48eb07f4814bd990386bbe94d32159dfd 100644 (file)
@@ -10,14 +10,16 @@ import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
-import static org.opendaylight.controller.cluster.datastore.messages.CreateTransaction.CURRENT_VERSION;
+import static org.opendaylight.controller.cluster.datastore.DataStoreVersions.CURRENT_VERSION;
 import akka.actor.ActorRef;
 import akka.actor.PoisonPill;
 import akka.actor.Props;
 import akka.dispatch.Dispatchers;
 import akka.dispatch.OnComplete;
 import akka.japi.Creator;
+import akka.japi.Procedure;
 import akka.pattern.Patterns;
+import akka.persistence.SnapshotSelectionCriteria;
 import akka.testkit.TestActorRef;
 import akka.util.Timeout;
 import com.google.common.base.Function;
@@ -43,6 +45,8 @@ import org.junit.Test;
 import org.mockito.InOrder;
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;
+import org.opendaylight.controller.cluster.DataPersistenceProvider;
+import org.opendaylight.controller.cluster.datastore.DatastoreContext.Builder;
 import org.opendaylight.controller.cluster.datastore.identifiers.ShardIdentifier;
 import org.opendaylight.controller.cluster.datastore.messages.AbortTransaction;
 import org.opendaylight.controller.cluster.datastore.messages.AbortTransactionReply;
@@ -59,12 +63,14 @@ import org.opendaylight.controller.cluster.datastore.messages.RegisterChangeList
 import org.opendaylight.controller.cluster.datastore.messages.UpdateSchemaContext;
 import org.opendaylight.controller.cluster.datastore.modification.MergeModification;
 import org.opendaylight.controller.cluster.datastore.modification.Modification;
+import org.opendaylight.controller.cluster.datastore.modification.ModificationPayload;
 import org.opendaylight.controller.cluster.datastore.modification.MutableCompositeModification;
 import org.opendaylight.controller.cluster.datastore.modification.WriteModification;
 import org.opendaylight.controller.cluster.datastore.node.NormalizedNodeToNodeCodec;
 import org.opendaylight.controller.cluster.datastore.utils.InMemoryJournal;
 import org.opendaylight.controller.cluster.datastore.utils.InMemorySnapshotStore;
 import org.opendaylight.controller.cluster.datastore.utils.MockDataChangeListener;
+import org.opendaylight.controller.cluster.datastore.utils.SerializationUtils;
 import org.opendaylight.controller.cluster.raft.ReplicatedLogEntry;
 import org.opendaylight.controller.cluster.raft.ReplicatedLogImplEntry;
 import org.opendaylight.controller.cluster.raft.Snapshot;
@@ -75,8 +81,8 @@ import org.opendaylight.controller.cluster.raft.base.messages.CaptureSnapshot;
 import org.opendaylight.controller.cluster.raft.base.messages.ElectionTimeout;
 import org.opendaylight.controller.cluster.raft.client.messages.FindLeader;
 import org.opendaylight.controller.cluster.raft.client.messages.FindLeaderReply;
+import org.opendaylight.controller.cluster.raft.protobuff.client.messages.CompositeModificationByteStringPayload;
 import org.opendaylight.controller.cluster.raft.protobuff.client.messages.CompositeModificationPayload;
-import org.opendaylight.controller.cluster.raft.protobuff.client.messages.Payload;
 import org.opendaylight.controller.md.cluster.datastore.model.SchemaContextHelper;
 import org.opendaylight.controller.md.cluster.datastore.model.TestModel;
 import org.opendaylight.controller.md.sal.common.api.data.AsyncDataBroker;
@@ -110,12 +116,13 @@ public class ShardTest extends AbstractActorTest {
     private final ShardIdentifier shardID = ShardIdentifier.builder().memberName("member-1")
             .shardName("inventory").type("config" + NEXT_SHARD_NUM.getAndIncrement()).build();
 
-    private DatastoreContext dataStoreContext = DatastoreContext.newBuilder().
+    private final Builder dataStoreContextBuilder = DatastoreContext.newBuilder().
             shardJournalRecoveryLogBatchSize(3).shardSnapshotBatchCount(5000).
-            shardHeartbeatIntervalInMillis(100).build();
+            shardHeartbeatIntervalInMillis(100);
 
     @Before
     public void setUp() {
+        Builder newBuilder = DatastoreContext.newBuilder();
         InMemorySnapshotStore.clear();
         InMemoryJournal.clear();
     }
@@ -126,9 +133,13 @@ public class ShardTest extends AbstractActorTest {
         InMemoryJournal.clear();
     }
 
+    private DatastoreContext newDatastoreContext() {
+        return dataStoreContextBuilder.build();
+    }
+
     private Props newShardProps() {
         return Shard.props(shardID, Collections.<ShardIdentifier,String>emptyMap(),
-                dataStoreContext, SCHEMA_CONTEXT);
+                newDatastoreContext(), SCHEMA_CONTEXT);
     }
 
     @Test
@@ -185,7 +196,7 @@ public class ShardTest extends AbstractActorTest {
                 @Override
                 public Shard create() throws Exception {
                     return new Shard(shardID, Collections.<ShardIdentifier,String>emptyMap(),
-                            dataStoreContext, SCHEMA_CONTEXT) {
+                            newDatastoreContext(), SCHEMA_CONTEXT) {
                         @Override
                         public void onReceiveCommand(final Object message) throws Exception {
                             if(message instanceof ElectionTimeout && firstElectionTimeout) {
@@ -313,7 +324,7 @@ public class ShardTest extends AbstractActorTest {
             class TestShard extends Shard {
                 TestShard() {
                     super(shardID, Collections.<ShardIdentifier, String>singletonMap(shardID, null),
-                            dataStoreContext, SCHEMA_CONTEXT);
+                            newDatastoreContext(), SCHEMA_CONTEXT);
                 }
 
                 Map<String, String> getPeerAddresses() {
@@ -357,13 +368,41 @@ public class ShardTest extends AbstractActorTest {
         TestActorRef<Shard> shard = TestActorRef.create(getSystem(), newShardProps(),
                 "testApplySnapshot");
 
-        NormalizedNodeToNodeCodec codec =
-            new NormalizedNodeToNodeCodec(SCHEMA_CONTEXT);
+        InMemoryDOMDataStore store = new InMemoryDOMDataStore("OPER", MoreExecutors.sameThreadExecutor());
+        store.onGlobalContextUpdated(SCHEMA_CONTEXT);
 
-        writeToStore(shard, TestModel.TEST_PATH, ImmutableNodes.containerNode(TestModel.TEST_QNAME));
+        writeToStore(store, TestModel.TEST_PATH, ImmutableNodes.containerNode(TestModel.TEST_QNAME));
 
         YangInstanceIdentifier root = YangInstanceIdentifier.builder().build();
-        NormalizedNode<?,?> expected = readStore(shard, root);
+        NormalizedNode<?,?> expected = readStore(store, root);
+
+        ApplySnapshot applySnapshot = new ApplySnapshot(Snapshot.create(
+                SerializationUtils.serializeNormalizedNode(expected),
+                Collections.<ReplicatedLogEntry>emptyList(), 1, 2, 3, 4));
+
+        shard.underlyingActor().onReceiveCommand(applySnapshot);
+
+        NormalizedNode<?,?> actual = readStore(shard, root);
+
+        assertEquals("Root node", expected, actual);
+
+        shard.tell(PoisonPill.getInstance(), ActorRef.noSender());
+    }
+
+    @Test
+    public void testApplyHelium2VersionSnapshot() throws Exception {
+        TestActorRef<Shard> shard = TestActorRef.create(getSystem(), newShardProps(),
+                "testApplySnapshot");
+
+        NormalizedNodeToNodeCodec codec = new NormalizedNodeToNodeCodec(SCHEMA_CONTEXT);
+
+        InMemoryDOMDataStore store = new InMemoryDOMDataStore("OPER", MoreExecutors.sameThreadExecutor());
+        store.onGlobalContextUpdated(SCHEMA_CONTEXT);
+
+        writeToStore(store, TestModel.TEST_PATH, ImmutableNodes.containerNode(TestModel.TEST_QNAME));
+
+        YangInstanceIdentifier root = YangInstanceIdentifier.builder().build();
+        NormalizedNode<?,?> expected = readStore(store, root);
 
         NormalizedNodeMessages.Container encode = codec.encode(expected);
 
@@ -375,7 +414,7 @@ public class ShardTest extends AbstractActorTest {
 
         NormalizedNode<?,?> actual = readStore(shard, root);
 
-        assertEquals(expected, actual);
+        assertEquals("Root node", expected, actual);
 
         shard.tell(PoisonPill.getInstance(), ActorRef.noSender());
     }
@@ -387,11 +426,26 @@ public class ShardTest extends AbstractActorTest {
 
         NormalizedNode<?, ?> node = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
 
-        MutableCompositeModification compMod = new MutableCompositeModification();
-        compMod.addModification(new WriteModification(TestModel.TEST_PATH, node, SCHEMA_CONTEXT));
-        Payload payload = new CompositeModificationPayload(compMod.toSerializable());
-        ApplyState applyState = new ApplyState(null, "test",
-                new ReplicatedLogImplEntry(1, 2, payload));
+        ApplyState applyState = new ApplyState(null, "test", new ReplicatedLogImplEntry(1, 2,
+                newModificationPayload(new WriteModification(TestModel.TEST_PATH, node))));
+
+        shard.underlyingActor().onReceiveCommand(applyState);
+
+        NormalizedNode<?,?> actual = readStore(shard, TestModel.TEST_PATH);
+        assertEquals("Applied state", node, actual);
+
+        shard.tell(PoisonPill.getInstance(), ActorRef.noSender());
+    }
+
+    @Test
+    public void testApplyStateLegacy() throws Exception {
+
+        TestActorRef<Shard> shard = TestActorRef.create(getSystem(), newShardProps(), "testApplyStateLegacy");
+
+        NormalizedNode<?, ?> node = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+        ApplyState applyState = new ApplyState(null, "test", new ReplicatedLogImplEntry(1, 2,
+                newLegacyByteStringPayload(new WriteModification(TestModel.TEST_PATH, node))));
 
         shard.underlyingActor().onReceiveCommand(applyState);
 
@@ -401,7 +455,6 @@ public class ShardTest extends AbstractActorTest {
         shard.tell(PoisonPill.getInstance(), ActorRef.noSender());
     }
 
-    @SuppressWarnings("serial")
     @Test
     public void testRecovery() throws Exception {
 
@@ -410,53 +463,107 @@ public class ShardTest extends AbstractActorTest {
         InMemoryDOMDataStore testStore = InMemoryDOMDataStoreFactory.create("Test", null, null);
         testStore.onGlobalContextUpdated(SCHEMA_CONTEXT);
 
-        DOMStoreWriteTransaction writeTx = testStore.newWriteOnlyTransaction();
-        writeTx.write(TestModel.TEST_PATH, ImmutableNodes.containerNode(TestModel.TEST_QNAME));
-        DOMStoreThreePhaseCommitCohort commitCohort = writeTx.ready();
-        commitCohort.preCommit().get();
-        commitCohort.commit().get();
+        writeToStore(testStore, TestModel.TEST_PATH, ImmutableNodes.containerNode(TestModel.TEST_QNAME));
 
-        DOMStoreReadTransaction readTx = testStore.newReadOnlyTransaction();
-        NormalizedNode<?, ?> root = readTx.read(YangInstanceIdentifier.builder().build()).get().get();
+        NormalizedNode<?, ?> root = readStore(testStore, YangInstanceIdentifier.builder().build());
 
         InMemorySnapshotStore.addSnapshot(shardID.toString(), Snapshot.create(
-                new NormalizedNodeToNodeCodec(SCHEMA_CONTEXT).encode(
-                        root).
-                                getNormalizedNode().toByteString().toByteArray(),
-                                Collections.<ReplicatedLogEntry>emptyList(), 0, 1, -1, -1));
+                SerializationUtils.serializeNormalizedNode(root),
+                Collections.<ReplicatedLogEntry>emptyList(), 0, 1, -1, -1));
 
         // Set up the InMemoryJournal.
 
-        InMemoryJournal.addEntry(shardID.toString(), 0, new ReplicatedLogImplEntry(0, 1, newPayload(
+        InMemoryJournal.addEntry(shardID.toString(), 0, new ReplicatedLogImplEntry(0, 1, newLegacyPayload(
                   new WriteModification(TestModel.OUTER_LIST_PATH,
-                          ImmutableNodes.mapNodeBuilder(TestModel.OUTER_LIST_QNAME).build(),
-                          SCHEMA_CONTEXT))));
+                          ImmutableNodes.mapNodeBuilder(TestModel.OUTER_LIST_QNAME).build()))));
 
-        int nListEntries = 11;
+        int nListEntries = 16;
         Set<Integer> listEntryKeys = new HashSet<>();
+
+        // Add some ModificationPayload entries
         for(int i = 1; i <= nListEntries; i++) {
             listEntryKeys.add(Integer.valueOf(i));
             YangInstanceIdentifier path = YangInstanceIdentifier.builder(TestModel.OUTER_LIST_PATH)
                     .nodeWithKey(TestModel.OUTER_LIST_QNAME, TestModel.ID_QNAME, i).build();
             Modification mod = new MergeModification(path,
-                    ImmutableNodes.mapEntry(TestModel.OUTER_LIST_QNAME, TestModel.ID_QNAME, i),
-                    SCHEMA_CONTEXT);
+                    ImmutableNodes.mapEntry(TestModel.OUTER_LIST_QNAME, TestModel.ID_QNAME, i));
             InMemoryJournal.addEntry(shardID.toString(), i, new ReplicatedLogImplEntry(i, 1,
-                    newPayload(mod)));
+                    newModificationPayload(mod)));
         }
 
         InMemoryJournal.addEntry(shardID.toString(), nListEntries + 1,
                 new ApplyLogEntries(nListEntries));
 
+        testRecovery(listEntryKeys);
+    }
+
+    @Test
+    public void testHelium2VersionRecovery() throws Exception {
+
+        // Set up the InMemorySnapshotStore.
+
+        InMemoryDOMDataStore testStore = InMemoryDOMDataStoreFactory.create("Test", null, null);
+        testStore.onGlobalContextUpdated(SCHEMA_CONTEXT);
+
+        writeToStore(testStore, TestModel.TEST_PATH, ImmutableNodes.containerNode(TestModel.TEST_QNAME));
+
+        NormalizedNode<?, ?> root = readStore(testStore, YangInstanceIdentifier.builder().build());
+
+        InMemorySnapshotStore.addSnapshot(shardID.toString(), Snapshot.create(
+                new NormalizedNodeToNodeCodec(SCHEMA_CONTEXT).encode(root).
+                                getNormalizedNode().toByteString().toByteArray(),
+                                Collections.<ReplicatedLogEntry>emptyList(), 0, 1, -1, -1));
+
+        // Set up the InMemoryJournal.
+
+        InMemoryJournal.addEntry(shardID.toString(), 0, new ReplicatedLogImplEntry(0, 1, newLegacyPayload(
+                  new WriteModification(TestModel.OUTER_LIST_PATH,
+                          ImmutableNodes.mapNodeBuilder(TestModel.OUTER_LIST_QNAME).build()))));
+
+        int nListEntries = 16;
+        Set<Integer> listEntryKeys = new HashSet<>();
+        int i = 1;
+
+        // Add some CompositeModificationPayload entries
+        for(; i <= 8; i++) {
+            listEntryKeys.add(Integer.valueOf(i));
+            YangInstanceIdentifier path = YangInstanceIdentifier.builder(TestModel.OUTER_LIST_PATH)
+                    .nodeWithKey(TestModel.OUTER_LIST_QNAME, TestModel.ID_QNAME, i).build();
+            Modification mod = new MergeModification(path,
+                    ImmutableNodes.mapEntry(TestModel.OUTER_LIST_QNAME, TestModel.ID_QNAME, i));
+            InMemoryJournal.addEntry(shardID.toString(), i, new ReplicatedLogImplEntry(i, 1,
+                    newLegacyPayload(mod)));
+        }
+
+        // Add some CompositeModificationByteStringPayload entries
+        for(; i <= nListEntries; i++) {
+            listEntryKeys.add(Integer.valueOf(i));
+            YangInstanceIdentifier path = YangInstanceIdentifier.builder(TestModel.OUTER_LIST_PATH)
+                    .nodeWithKey(TestModel.OUTER_LIST_QNAME, TestModel.ID_QNAME, i).build();
+            Modification mod = new MergeModification(path,
+                    ImmutableNodes.mapEntry(TestModel.OUTER_LIST_QNAME, TestModel.ID_QNAME, i));
+            InMemoryJournal.addEntry(shardID.toString(), i, new ReplicatedLogImplEntry(i, 1,
+                    newLegacyByteStringPayload(mod)));
+        }
+
+        InMemoryJournal.addEntry(shardID.toString(), nListEntries + 1, new ApplyLogEntries(nListEntries));
+
+        testRecovery(listEntryKeys);
+    }
+
+    private void testRecovery(Set<Integer> listEntryKeys) throws Exception {
         // Create the actor and wait for recovery complete.
 
+        int nListEntries = listEntryKeys.size();
+
         final CountDownLatch recoveryComplete = new CountDownLatch(1);
 
+        @SuppressWarnings("serial")
         Creator<Shard> creator = new Creator<Shard>() {
             @Override
             public Shard create() throws Exception {
                 return new Shard(shardID, Collections.<ShardIdentifier,String>emptyMap(),
-                        dataStoreContext, SCHEMA_CONTEXT) {
+                        newDatastoreContext(), SCHEMA_CONTEXT) {
                     @Override
                     protected void onRecoveryComplete() {
                         try {
@@ -507,7 +614,7 @@ public class ShardTest extends AbstractActorTest {
         shard.tell(PoisonPill.getInstance(), ActorRef.noSender());
     }
 
-    private CompositeModificationPayload newPayload(final Modification... mods) {
+    private CompositeModificationPayload newLegacyPayload(final Modification... mods) {
         MutableCompositeModification compMod = new MutableCompositeModification();
         for(Modification mod: mods) {
             compMod.addModification(mod);
@@ -516,6 +623,24 @@ public class ShardTest extends AbstractActorTest {
         return new CompositeModificationPayload(compMod.toSerializable());
     }
 
+    private CompositeModificationByteStringPayload newLegacyByteStringPayload(final Modification... mods) {
+        MutableCompositeModification compMod = new MutableCompositeModification();
+        for(Modification mod: mods) {
+            compMod.addModification(mod);
+        }
+
+        return new CompositeModificationByteStringPayload(compMod.toSerializable());
+    }
+
+    private ModificationPayload newModificationPayload(final Modification... mods) throws IOException {
+        MutableCompositeModification compMod = new MutableCompositeModification();
+        for(Modification mod: mods) {
+            compMod.addModification(mod);
+        }
+
+        return new ModificationPayload(compMod);
+    }
+
     private DOMStoreThreePhaseCommitCohort setupMockWriteTransaction(final String cohortName,
             final InMemoryDOMDataStore dataStore, final YangInstanceIdentifier path, final NormalizedNode<?, ?> data,
             final MutableCompositeModification modification) {
@@ -564,7 +689,7 @@ public class ShardTest extends AbstractActorTest {
             }
         }).when(cohort).abort();
 
-        modification.addModification(new WriteModification(path, data, SCHEMA_CONTEXT));
+        modification.addModification(new WriteModification(path, data));
 
         return cohort;
     }
@@ -750,15 +875,72 @@ public class ShardTest extends AbstractActorTest {
             assertTrue("Missing leaf " + TestModel.ID_QNAME.getLocalName(), idLeaf.isPresent());
             assertEquals(TestModel.ID_QNAME.getLocalName() + " value", 1, idLeaf.get().getValue());
 
-            for(int i = 0; i < 20 * 5; i++) {
-                long lastLogIndex = shard.underlyingActor().getShardMBean().getLastLogIndex();
-                if(lastLogIndex == 2) {
-                    break;
-                }
-                Uninterruptibles.sleepUninterruptibly(50, TimeUnit.MILLISECONDS);
+            verifyLastLogIndex(shard, 2);
+
+            shard.tell(PoisonPill.getInstance(), ActorRef.noSender());
+        }};
+    }
+
+    private void verifyLastLogIndex(TestActorRef<Shard> shard, long expectedValue) {
+        for(int i = 0; i < 20 * 5; i++) {
+            long lastLogIndex = shard.underlyingActor().getShardMBean().getLastLogIndex();
+            if(lastLogIndex == expectedValue) {
+                break;
             }
+            Uninterruptibles.sleepUninterruptibly(50, TimeUnit.MILLISECONDS);
+        }
+
+        assertEquals("Last log index", expectedValue, shard.underlyingActor().getShardMBean().getLastLogIndex());
+    }
+
+    @Test
+    public void testCommitWithPersistenceDisabled() throws Throwable {
+        dataStoreContextBuilder.persistent(false);
+        new ShardTestKit(getSystem()) {{
+            final TestActorRef<Shard> shard = TestActorRef.create(getSystem(),
+                    newShardProps().withDispatcher(Dispatchers.DefaultDispatcherId()),
+                    "testCommitPhaseFailure");
+
+            waitUntilLeader(shard);
+
+            InMemoryDOMDataStore dataStore = shard.underlyingActor().getDataStore();
+
+            // Setup a simulated transactions with a mock cohort.
+
+            String transactionID = "tx";
+            MutableCompositeModification modification = new MutableCompositeModification();
+            NormalizedNode<?, ?> containerNode = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+            DOMStoreThreePhaseCommitCohort cohort = setupMockWriteTransaction("cohort", dataStore,
+                    TestModel.TEST_PATH, containerNode, modification);
+
+            FiniteDuration duration = duration("5 seconds");
+
+            // Simulate the ForwardedReadyTransaction messages that would be sent
+            // by the ShardTransaction.
+
+            shard.tell(new ForwardedReadyTransaction(transactionID, CURRENT_VERSION,
+                    cohort, modification, true), getRef());
+            expectMsgClass(duration, ReadyTransactionReply.SERIALIZABLE_CLASS);
+
+            // Send the CanCommitTransaction message.
+
+            shard.tell(new CanCommitTransaction(transactionID).toSerializable(), getRef());
+            CanCommitTransactionReply canCommitReply = CanCommitTransactionReply.fromSerializable(
+                    expectMsgClass(duration, CanCommitTransactionReply.SERIALIZABLE_CLASS));
+            assertEquals("Can commit", true, canCommitReply.getCanCommit());
+
+            // Send the CanCommitTransaction message.
+
+            shard.tell(new CommitTransaction(transactionID).toSerializable(), getRef());
+            expectMsgClass(duration, CommitTransactionReply.SERIALIZABLE_CLASS);
+
+            InOrder inOrder = inOrder(cohort);
+            inOrder.verify(cohort).canCommit();
+            inOrder.verify(cohort).preCommit();
+            inOrder.verify(cohort).commit();
 
-            assertEquals("Last log index", 2, shard.underlyingActor().getShardMBean().getLastLogIndex());
+            NormalizedNode<?, ?> actualNode = readStore(shard, TestModel.TEST_PATH);
+            assertEquals(TestModel.TEST_QNAME.getLocalName(), containerNode, actualNode);
 
             shard.tell(PoisonPill.getInstance(), ActorRef.noSender());
         }};
@@ -930,26 +1112,24 @@ public class ShardTest extends AbstractActorTest {
             waitUntilLeader(shard);
 
             final FiniteDuration duration = duration("5 seconds");
-            final Timeout timeout = new Timeout(duration);
-
             InMemoryDOMDataStore dataStore = shard.underlyingActor().getDataStore();
 
             final String transactionID = "tx1";
-            final CountDownLatch abortComplete = new CountDownLatch(1);
             Function<DOMStoreThreePhaseCommitCohort,ListenableFuture<Void>> preCommit =
                           new Function<DOMStoreThreePhaseCommitCohort,ListenableFuture<Void>>() {
                 @Override
                 public ListenableFuture<Void> apply(final DOMStoreThreePhaseCommitCohort cohort) {
                     ListenableFuture<Void> preCommitFuture = cohort.preCommit();
 
-                    Future<Object> abortFuture = Patterns.ask(shard,
-                            new AbortTransaction(transactionID).toSerializable(), timeout);
-                    abortFuture.onComplete(new OnComplete<Object>() {
-                        @Override
-                        public void onComplete(final Throwable e, final Object resp) {
-                            abortComplete.countDown();
-                        }
-                    }, getSystem().dispatcher());
+                    // Simulate an AbortTransaction message occurring during replication, after
+                    // persisting and before finishing the commit to the in-memory store.
+                    // We have no followers so due to optimizations in the RaftActor, it does not
+                    // attempt replication and thus we can't send an AbortTransaction message b/c
+                    // it would be processed too late after CommitTransaction completes. So we'll
+                    // simulate an AbortTransaction message occurring during replication by calling
+                    // the shard directly.
+                    //
+                    shard.underlyingActor().doAbortTransaction(transactionID, null);
 
                     return preCommitFuture;
                 }
@@ -969,14 +1149,14 @@ public class ShardTest extends AbstractActorTest {
                     expectMsgClass(duration, CanCommitTransactionReply.SERIALIZABLE_CLASS));
             assertEquals("Can commit", true, canCommitReply.getCanCommit());
 
-            Future<Object> commitFuture = Patterns.ask(shard,
-                    new CommitTransaction(transactionID).toSerializable(), timeout);
-
-            assertEquals("Abort complete", true, abortComplete.await(5, TimeUnit.SECONDS));
-
-            Await.result(commitFuture, duration);
+            shard.tell(new CommitTransaction(transactionID).toSerializable(), getRef());
+            expectMsgClass(duration, CommitTransactionReply.SERIALIZABLE_CLASS);
 
             NormalizedNode<?, ?> node = readStore(shard, TestModel.TEST_PATH);
+
+            // Since we're simulating an abort occurring during replication and before finish commit,
+            // the data should still get written to the in-memory store since we've gotten past
+            // canCommit and preCommit and persisted the data.
             assertNotNull(TestModel.TEST_QNAME.getLocalName() + " not found", node);
 
             shard.tell(PoisonPill.getInstance(), ActorRef.noSender());
@@ -985,7 +1165,7 @@ public class ShardTest extends AbstractActorTest {
 
     @Test
     public void testTransactionCommitTimeout() throws Throwable {
-        dataStoreContext = DatastoreContext.newBuilder().shardTransactionCommitTimeoutInSeconds(1).build();
+        dataStoreContextBuilder.shardTransactionCommitTimeoutInSeconds(1);
 
         new ShardTestKit(getSystem()) {{
             final TestActorRef<Shard> shard = TestActorRef.create(getSystem(),
@@ -1057,7 +1237,7 @@ public class ShardTest extends AbstractActorTest {
 
     @Test
     public void testTransactionCommitQueueCapacityExceeded() throws Throwable {
-        dataStoreContext = DatastoreContext.newBuilder().shardTransactionCommitQueueCapacity(1).build();
+        dataStoreContextBuilder.shardTransactionCommitQueueCapacity(1);
 
         new ShardTestKit(getSystem()) {{
             final TestActorRef<Shard> shard = TestActorRef.create(getSystem(),
@@ -1190,15 +1370,7 @@ public class ShardTest extends AbstractActorTest {
 
             // Wait for the 2nd Tx to complete the canCommit phase.
 
-            final CountDownLatch latch = new CountDownLatch(1);
-            canCommitFuture.onComplete(new OnComplete<Object>() {
-                @Override
-                public void onComplete(final Throwable t, final Object resp) {
-                    latch.countDown();
-                }
-            }, getSystem().dispatcher());
-
-            assertEquals("2nd CanCommit complete", true, latch.await(5, TimeUnit.SECONDS));
+            Await.ready(canCommitFuture, duration);
 
             InOrder inOrder = inOrder(cohort1, cohort2);
             inOrder.verify(cohort1).canCommit();
@@ -1209,19 +1381,54 @@ public class ShardTest extends AbstractActorTest {
     }
 
     @Test
-    public void testCreateSnapshot() throws IOException, InterruptedException {
-            testCreateSnapshot(true, "testCreateSnapshot");
+    public void testCreateSnapshot() throws Exception {
+        testCreateSnapshot(true, "testCreateSnapshot");
     }
 
     @Test
-    public void testCreateSnapshotWithNonPersistentData() throws IOException, InterruptedException {
+    public void testCreateSnapshotWithNonPersistentData() throws Exception {
         testCreateSnapshot(false, "testCreateSnapshotWithNonPersistentData");
     }
 
     @SuppressWarnings("serial")
-    public void testCreateSnapshot(final boolean persistent, final String shardActorName) throws IOException, InterruptedException {
-        final DatastoreContext dataStoreContext = DatastoreContext.newBuilder().
-                shardJournalRecoveryLogBatchSize(3).shardSnapshotBatchCount(5000).persistent(persistent).build();
+    public void testCreateSnapshot(final boolean persistent, final String shardActorName) throws Exception{
+
+        final AtomicReference<Object> savedSnapshot = new AtomicReference<>();
+        class DelegatingPersistentDataProvider implements DataPersistenceProvider {
+            DataPersistenceProvider delegate;
+
+            DelegatingPersistentDataProvider(DataPersistenceProvider delegate) {
+                this.delegate = delegate;
+            }
+
+            @Override
+            public boolean isRecoveryApplicable() {
+                return delegate.isRecoveryApplicable();
+            }
+
+            @Override
+            public <T> void persist(T o, Procedure<T> procedure) {
+                delegate.persist(o, procedure);
+            }
+
+            @Override
+            public void saveSnapshot(Object o) {
+                savedSnapshot.set(o);
+                delegate.saveSnapshot(o);
+            }
+
+            @Override
+            public void deleteSnapshots(SnapshotSelectionCriteria criteria) {
+                delegate.deleteSnapshots(criteria);
+            }
+
+            @Override
+            public void deleteMessages(long sequenceNumber) {
+                delegate.deleteMessages(sequenceNumber);
+            }
+        }
+
+        dataStoreContextBuilder.persistent(persistent);
 
         new ShardTestKit(getSystem()) {{
             final AtomicReference<CountDownLatch> latch = new AtomicReference<>(new CountDownLatch(1));
@@ -1229,7 +1436,19 @@ public class ShardTest extends AbstractActorTest {
                 @Override
                 public Shard create() throws Exception {
                     return new Shard(shardID, Collections.<ShardIdentifier,String>emptyMap(),
-                            dataStoreContext, SCHEMA_CONTEXT) {
+                            newDatastoreContext(), SCHEMA_CONTEXT) {
+
+                        DelegatingPersistentDataProvider delegating;
+
+                        @Override
+                        protected DataPersistenceProvider persistence() {
+                            if(delegating == null) {
+                                delegating = new DelegatingPersistentDataProvider(super.persistence());
+                            }
+
+                            return delegating;
+                        }
+
                         @Override
                         protected void commitSnapshot(final long sequenceNumber) {
                             super.commitSnapshot(sequenceNumber);
@@ -1244,16 +1463,40 @@ public class ShardTest extends AbstractActorTest {
 
             waitUntilLeader(shard);
 
-            shard.tell(new CaptureSnapshot(-1,-1,-1,-1), getRef());
+            writeToStore(shard, TestModel.TEST_PATH, ImmutableNodes.containerNode(TestModel.TEST_QNAME));
+
+            NormalizedNode<?,?> expectedRoot = readStore(shard, YangInstanceIdentifier.builder().build());
+
+            CaptureSnapshot capture = new CaptureSnapshot(-1, -1, -1, -1, -1, -1);
+            shard.tell(capture, getRef());
 
             assertEquals("Snapshot saved", true, latch.get().await(5, TimeUnit.SECONDS));
 
+            assertTrue("Invalid saved snapshot " + savedSnapshot.get(),
+                    savedSnapshot.get() instanceof Snapshot);
+
+            verifySnapshot((Snapshot)savedSnapshot.get(), expectedRoot);
+
             latch.set(new CountDownLatch(1));
-            shard.tell(new CaptureSnapshot(-1,-1,-1,-1), getRef());
+            savedSnapshot.set(null);
+
+            shard.tell(capture, getRef());
 
             assertEquals("Snapshot saved", true, latch.get().await(5, TimeUnit.SECONDS));
 
+            assertTrue("Invalid saved snapshot " + savedSnapshot.get(),
+                    savedSnapshot.get() instanceof Snapshot);
+
+            verifySnapshot((Snapshot)savedSnapshot.get(), expectedRoot);
+
             shard.tell(PoisonPill.getInstance(), ActorRef.noSender());
+        }
+
+        private void verifySnapshot(Snapshot snapshot, NormalizedNode<?,?> expectedRoot) {
+
+            NormalizedNode<?, ?> actual = SerializationUtils.deserializeNormalizedNode(snapshot.getState());
+            assertEquals("Root node", expectedRoot, actual);
+
         }};
     }
 
@@ -1360,7 +1603,12 @@ public class ShardTest extends AbstractActorTest {
 
     static NormalizedNode<?,?> readStore(final TestActorRef<Shard> shard, final YangInstanceIdentifier id)
             throws ExecutionException, InterruptedException {
-        DOMStoreReadTransaction transaction = shard.underlyingActor().getDataStore().newReadOnlyTransaction();
+        return readStore(shard.underlyingActor().getDataStore(), id);
+    }
+
+    public static NormalizedNode<?,?> readStore(final InMemoryDOMDataStore store, final YangInstanceIdentifier id)
+            throws ExecutionException, InterruptedException {
+        DOMStoreReadTransaction transaction = store.newReadOnlyTransaction();
 
         CheckedFuture<Optional<NormalizedNode<?, ?>>, ReadFailedException> future =
             transaction.read(id);
@@ -1373,9 +1621,14 @@ public class ShardTest extends AbstractActorTest {
         return node;
     }
 
-    private void writeToStore(final TestActorRef<Shard> shard, final YangInstanceIdentifier id, final NormalizedNode<?,?> node)
-        throws ExecutionException, InterruptedException {
-        DOMStoreWriteTransaction transaction = shard.underlyingActor().getDataStore().newWriteOnlyTransaction();
+    static void writeToStore(final TestActorRef<Shard> shard, final YangInstanceIdentifier id,
+            final NormalizedNode<?,?> node) throws ExecutionException, InterruptedException {
+        writeToStore(shard.underlyingActor().getDataStore(), id, node);
+    }
+
+    public static void writeToStore(final InMemoryDOMDataStore store, final YangInstanceIdentifier id,
+            final NormalizedNode<?,?> node) throws ExecutionException, InterruptedException {
+        DOMStoreWriteTransaction transaction = store.newWriteOnlyTransaction();
 
         transaction.write(id, node);