X-Git-Url: https://git.opendaylight.org/gerrit/gitweb?p=controller.git;a=blobdiff_plain;f=opendaylight%2Fmd-sal%2Fsal-akka-raft%2Fsrc%2Ftest%2Fjava%2Forg%2Fopendaylight%2Fcontroller%2Fcluster%2Fraft%2Futils%2FInMemorySnapshotStore.java;h=e16d7949745004ca91dbf20e4d186beb8d84cf50;hp=01f337567560ba02f7849c8b920cd5f05b18f949;hb=11dadddb4d9ba26ae0b1795921c7a218a6d893c2;hpb=753515e8868a1a15982d3f2697439f522f273db5

diff --git a/opendaylight/md-sal/sal-akka-raft/src/test/java/org/opendaylight/controller/cluster/raft/utils/InMemorySnapshotStore.java b/opendaylight/md-sal/sal-akka-raft/src/test/java/org/opendaylight/controller/cluster/raft/utils/InMemorySnapshotStore.java
index 01f3375675..e16d794974 100644
--- a/opendaylight/md-sal/sal-akka-raft/src/test/java/org/opendaylight/controller/cluster/raft/utils/InMemorySnapshotStore.java
+++ b/opendaylight/md-sal/sal-akka-raft/src/test/java/org/opendaylight/controller/cluster/raft/utils/InMemorySnapshotStore.java
@@ -16,12 +16,15 @@ import akka.persistence.SnapshotSelectionCriteria;
 import akka.persistence.snapshot.japi.SnapshotStore;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
+import com.google.common.util.concurrent.Uninterruptibles;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import scala.concurrent.Future;
@@ -36,6 +39,7 @@ public class InMemorySnapshotStore extends SnapshotStore {
     static final Logger LOG = LoggerFactory.getLogger(InMemorySnapshotStore.class);
 
     private static Map<String, List<StoredSnapshot>> snapshots = new ConcurrentHashMap<>();
+    private static final Map<String, CountDownLatch> snapshotSavedLatches = new ConcurrentHashMap<>();
 
     public static void addSnapshot(String persistentId, Object snapshot) {
         List<StoredSnapshot> snapshotList = snapshots.get(persistentId);
@@ -75,6 +79,18 @@ public class InMemorySnapshotStore extends SnapshotStore {
         snapshots.clear();
     }
 
+    public static void addSnapshotSavedLatch(String persistenceId) {
+        snapshotSavedLatches.put(persistenceId, new CountDownLatch(1));
+    }
+
+    public static <T> T waitForSavedSnapshot(String persistenceId, Class<T> type) {
+        if(!Uninterruptibles.awaitUninterruptibly(snapshotSavedLatches.get(persistenceId), 5, TimeUnit.SECONDS)) {
+            throw new AssertionError("Snapshot was not saved");
+        }
+
+        return getSnapshots(persistenceId, type).get(0);
+    }
+
     @Override
     public Future<Option<SelectedSnapshot>> doLoadAsync(String s,
         SnapshotSelectionCriteria snapshotSelectionCriteria) {
@@ -101,6 +117,11 @@ public class InMemorySnapshotStore extends SnapshotStore {
             snapshotList.add(new StoredSnapshot(snapshotMetadata, o));
         }
 
+        CountDownLatch latch = snapshotSavedLatches.get(snapshotMetadata.persistenceId());
+        if(latch != null) {
+            latch.countDown();
+        }
+
         return Futures.successful(null);
     }