X-Git-Url: https://git.opendaylight.org/gerrit/gitweb?a=blobdiff_plain;f=opendaylight%2Fmd-sal%2Fsal-distributed-datastore%2Fsrc%2Ftest%2Fjava%2Forg%2Fopendaylight%2Fcontroller%2Fcluster%2Fdatastore%2FTestShard.java;h=9eb20c03960792898a611d89736a6d4a3b205031;hb=229c55fbdc5907129c9c35b3e98963b5fa2528f0;hp=b0e744a24add221131082faba9e3144f8b35c497;hpb=3115b8171461584e85f58d87a9f179013cfbb262;p=controller.git

diff --git a/opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/TestShard.java b/opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/TestShard.java
index b0e744a24a..9eb20c0396 100644
--- a/opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/TestShard.java
+++ b/opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/TestShard.java
@@ -7,14 +7,51 @@
  */
 package org.opendaylight.controller.cluster.datastore;
 
+import static java.util.Objects.requireNonNull;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Predicate;
 import org.opendaylight.controller.cluster.datastore.persisted.FrontendShardDataTreeSnapshotMetadata;
 
 public class TestShard extends Shard {
+    public static class Builder extends Shard.Builder {
+        Builder() {
+            super(TestShard.class);
+        }
+    }
+
     // Message to request FrontendMetadata
     public static final class RequestFrontendMetadata {
 
     }
 
+    private abstract static class DropMessages<T> {
+        private final Class<T> msgClass;
+
+        DropMessages(final Class<T> msgClass) {
+            this.msgClass = requireNonNull(msgClass);
+        }
+
+        final Class<T> getMsgClass() {
+            return msgClass;
+        }
+    }
+
+    public static class StartDropMessages<T> extends DropMessages<T> {
+        public StartDropMessages(final Class<T> msgClass) {
+            super(msgClass);
+        }
+    }
+
+    public static class StopDropMessages<T> extends DropMessages<T> {
+        public StopDropMessages(final Class<T> msgClass) {
+            super(msgClass);
+        }
+    }
+
+    private final Map<Class<?>, Predicate<?>> dropMessages = new ConcurrentHashMap<>();
+
     protected TestShard(AbstractBuilder<?, ?> builder) {
         super(builder);
     }
@@ -29,14 +66,37 @@ public class TestShard extends Shard {
         }
     }
 
-    public static Shard.Builder builder() {
-        return new TestShard.Builder();
+    @Override
+    protected void handleCommand(Object message) {
+        if (message instanceof StartDropMessages) {
+            startDropMessages(((StartDropMessages<?>) message).getMsgClass());
+        } else if (message instanceof StopDropMessages) {
+            stopDropMessages(((StopDropMessages<?>) message).getMsgClass());
+        } else {
+            dropOrHandle(message);
+        }
     }
 
-    public static class Builder extends Shard.Builder {
-        Builder() {
-            super(TestShard.class);
+    private <T> void dropOrHandle(T message) {
+        Predicate<T> drop = (Predicate<T>) dropMessages.get(message.getClass());
+        if (drop == null || !drop.test(message)) {
+            super.handleCommand(message);
         }
     }
-}
 
+    private void startDropMessages(final Class<?> msgClass) {
+        dropMessages.put(msgClass, msg -> true);
+    }
+
+    <T> void startDropMessages(final Class<T> msgClass, final Predicate<T> filter) {
+        dropMessages.put(msgClass, filter);
+    }
+
+    public void stopDropMessages(final Class<?> msgClass) {
+        dropMessages.remove(msgClass);
+    }
+
+    public static TestShard.Builder builder() {
+        return new TestShard.Builder();
+    }
+}