SOLR-14248: Improve ClusterStateMockUtil and make its methods public

solr/CHANGES.txt

@@ -297,6 +297,8 @@ Other Changes
   TestInjection.skipIndexWriterCommitOnClose.  Users that modified DUH2.commitOnClose in test cases for custom
   plugins/modicitations should now use TestInjection instead.  (hossman)
 
+* SOLR-14248: Improve ClusterStateMockUtil and make its methods public. (shalin)
+
 ==================  8.4.1 ==================
 
 Consult the LUCENE_CHANGES.txt file for additional, low level, changes in this release.

solr/core/src/test/org/apache/solr/cloud/ClusterStateMockUtil.java

@@ -27,27 +27,35 @@ import java.util.regex.Pattern;
 
 import org.apache.solr.common.cloud.ClusterState;
 import org.apache.solr.common.cloud.DocCollection;
+import org.apache.solr.common.cloud.DocRouter;
 import org.apache.solr.common.cloud.Replica;
 import org.apache.solr.common.cloud.Slice;
 import org.apache.solr.common.cloud.ZkStateReader;
 import org.apache.solr.common.util.Utils;
 
+/**
+ * A utility class that can create mock ZkStateReader objects with custom ClusterState objects created
+ * using a simple string based description. See {@link #buildClusterState(String, int, int, String...)} for
+ * details on how the cluster state can be created.
+ *
+ * @lucene.experimental
+ */
 public class ClusterStateMockUtil {
 
   private final static Pattern BLUEPRINT = Pattern.compile("([a-z])(\\d+)?(?:(['A','R','D','F']))?(\\*)?");
 
-  protected static ZkStateReader buildClusterState(String string, String ... liveNodes) {
+  public static ZkStateReader buildClusterState(String clusterDescription, String ... liveNodes) {
-    return buildClusterState(string, 1, liveNodes);
+    return buildClusterState(clusterDescription, 1, liveNodes);
   }
 
-  protected static ZkStateReader buildClusterState(String string, int replicationFactor, String ... liveNodes) {
+  public static ZkStateReader buildClusterState(String clusterDescription, int replicationFactor, String ... liveNodes) {
-    return buildClusterState(string, replicationFactor, 10, liveNodes);
+    return buildClusterState(clusterDescription, replicationFactor, 10, liveNodes);
   }
 
   /**
    * This method lets you construct a complex ClusterState object by using simple strings of letters.
    *
-   * c = collection, s = slice, r = replica, \d = node number (r2 means the replica is on node 2),
+   * c = collection, s = slice, r = replica (nrt type, default), n = nrt replica, t = tlog replica, p = pull replica, \d = node number (r2 means the replica is on node 2),
    * state = [A,R,D,F], * = replica to replace, binds to the left.
    *
    * For example:
@@ -104,7 +112,7 @@ public class ClusterStateMockUtil {
    *
    */
   @SuppressWarnings("resource")
-  protected static ZkStateReader buildClusterState(String clusterDescription, int replicationFactor, int maxShardsPerNode, String ... liveNodes) {
+  public static ZkStateReader buildClusterState(String clusterDescription, int replicationFactor, int maxShardsPerNode, String ... liveNodes) {
     Map<String,Slice> slices = null;
     Map<String,Replica> replicas = null;
     Map<String,Object> collectionProps = new HashMap<>();
@@ -123,7 +131,7 @@ public class ClusterStateMockUtil {
       switch (m.group(1)) {
         case "c":
           slices = new HashMap<>();
-          docCollection = new DocCollection(collName = "collection" + (collectionStates.size() + 1), slices, collectionProps, null);
+          docCollection = new DocCollection(collName = "collection" + (collectionStates.size() + 1), slices, collectionProps, DocRouter.DEFAULT);
           collectionStates.put(docCollection.getName(), docCollection);
           break;
         case "s":
@@ -131,49 +139,25 @@ public class ClusterStateMockUtil {
           if(collName == null) collName = "collection" + (collectionStates.size() + 1);
           slice = new Slice(sliceName = "slice" + (slices.size() + 1), replicas, null,  collName);
           slices.put(slice.getName(), slice);
+
+          // hack alert: the DocCollection constructor copies over active slices to its active slice map in the constructor
+          // but here we construct the DocCollection before creating the slices which breaks code that calls DocCollection.getActiveSlices
+          // so here we re-create doc collection with the latest slices map to workaround this problem
+          // todo: a better fix would be to have a builder class for DocCollection that builds the final object once all the slices and replicas have been created.
+          docCollection = docCollection.copyWithSlices(slices);
+          collectionStates.put(docCollection.getName(), docCollection);
           break;
         case "r":
-          Map<String,Object> replicaPropMap = new HashMap<>();
+        case "n":
-          String node;
+        case "t":
-
+        case "p":
-          node = m.group(2);
+          String node = m.group(2);
-
+          String replicaName = "replica" + replicaCount++;
-          if (node == null || node.trim().length() == 0) {
-            node = "1";
-          }
-
-          Replica.State state = Replica.State.ACTIVE;
           String stateCode = m.group(3);
 
-          if (stateCode != null) {
+          Map<String, Object> replicaPropMap = makeReplicaProps(sliceName, node, replicaName, stateCode, m.group(1));
-            switch (stateCode.charAt(0)) {
+          if (collName == null) collName = "collection" + (collectionStates.size() + 1);
-              case 'S':
+          if (sliceName == null) collName = "slice" + (slices.size() + 1);
-                state = Replica.State.ACTIVE;
-                break;
-              case 'R':
-                state = Replica.State.RECOVERING;
-                break;
-              case 'D':
-                state = Replica.State.DOWN;
-                break;
-              case 'F':
-                state = Replica.State.RECOVERY_FAILED;
-                break;
-              default:
-                throw new IllegalArgumentException(
-                    "Unexpected state for replica: " + stateCode);
-            }
-          }
-
-          String nodeName = "baseUrl" + node + "_";
-          String replicaName = "replica" + replicaCount++;
-
-          replicaPropMap.put(ZkStateReader.NODE_NAME_PROP, nodeName);
-          replicaPropMap.put(ZkStateReader.BASE_URL_PROP, "http://baseUrl" + node);
-          replicaPropMap.put(ZkStateReader.STATE_PROP, state.toString());
-          replicaPropMap.put(ZkStateReader.CORE_NAME_PROP, "core_" + replicaName);
-          if(collName == null) collName = "collection" + (collectionStates.size() + 1);
-          if(sliceName == null) collName = "slice" + (slices.size() + 1);
           replica = new Replica(replicaName, replicaPropMap, collName, sliceName);
 
           replicas.put(replica.getName(), replica);
@@ -193,5 +177,50 @@ public class ClusterStateMockUtil {
     return reader;
   }
 
+  private static Map<String, Object> makeReplicaProps(String sliceName, String node, String replicaName, String stateCode, String replicaTypeCode) {
+    if (node == null || node.trim().length() == 0) {
+      node = "1";
+    }
+
+    Replica.State state = Replica.State.ACTIVE;
+    if (stateCode != null) {
+      switch (stateCode.charAt(0)) {
+        case 'S':
+          state = Replica.State.ACTIVE;
+          break;
+        case 'R':
+          state = Replica.State.RECOVERING;
+          break;
+        case 'D':
+          state = Replica.State.DOWN;
+          break;
+        case 'F':
+          state = Replica.State.RECOVERY_FAILED;
+          break;
+        default:
+          throw new IllegalArgumentException(
+              "Unexpected state for replica: " + stateCode);
+      }
+    }
+
+    Replica.Type replicaType = Replica.Type.NRT;
+    switch (replicaTypeCode)  {
+      case "t":
+        replicaType = Replica.Type.TLOG;
+        break;
+      case "p":
+        replicaType = Replica.Type.PULL;
+        break;
+    }
+
+    Map<String,Object> replicaPropMap = new HashMap<>();
+    replicaPropMap.put(ZkStateReader.NODE_NAME_PROP, "baseUrl" + node + "_");
+    replicaPropMap.put(ZkStateReader.BASE_URL_PROP, "http://baseUrl" + node);
+    replicaPropMap.put(ZkStateReader.STATE_PROP, state.toString());
+    replicaPropMap.put(ZkStateReader.CORE_NAME_PROP, sliceName + "_" + replicaName);
+    replicaPropMap.put(ZkStateReader.REPLICA_TYPE, replicaType.name());
+    return replicaPropMap;
+  }
+
 
 }

solr/core/src/test/org/apache/solr/cloud/ClusterStateMockUtilTest.java

@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.solr.cloud;
+
+import org.apache.solr.SolrTestCaseJ4;
+import org.apache.solr.common.cloud.ClusterState;
+import org.apache.solr.common.cloud.DocCollection;
+import org.apache.solr.common.cloud.DocRouter;
+import org.apache.solr.common.cloud.Replica;
+import org.apache.solr.common.cloud.Slice;
+import org.apache.solr.common.cloud.ZkStateReader;
+import org.junit.Test;
+
+/**
+ * Tests for {@link ClusterStateMockUtil}
+ */
+public class ClusterStateMockUtilTest extends SolrTestCaseJ4 {
+
+  @Test
+  public void testBuildClusterState_Simple() {
+    try (ZkStateReader zkStateReader = ClusterStateMockUtil.buildClusterState("csr", "baseUrl1_")) {
+      ClusterState clusterState = zkStateReader.getClusterState();
+      assertNotNull(clusterState);
+      assertEquals(1, clusterState.getCollectionStates().size());
+      DocCollection collection1 = clusterState.getCollectionOrNull("collection1");
+      assertNotNull(collection1);
+      assertEquals(DocRouter.DEFAULT, collection1.getRouter());
+      assertEquals(1, collection1.getActiveSlices().size());
+      assertEquals(1, collection1.getSlices().size());
+      Slice slice1 = collection1.getSlice("slice1");
+      assertNotNull(slice1);
+      assertEquals(1, slice1.getReplicas().size());
+      Replica replica1 = slice1.getReplica("replica1");
+      assertNotNull(replica1);
+      assertEquals("baseUrl1_", replica1.getNodeName());
+      assertEquals("slice1_replica1", replica1.getCoreName());
+      assertEquals("http://baseUrl1", replica1.getBaseUrl());
+      assertEquals("http://baseUrl1/slice1_replica1/", replica1.getCoreUrl());
+      assertEquals(Replica.State.ACTIVE, replica1.getState());
+      assertEquals(Replica.Type.NRT, replica1.getType());
+    }
+  }
+
+  @Test
+  public void testBuildClusterState_ReplicaTypes() {
+    try (ZkStateReader zkStateReader = ClusterStateMockUtil.buildClusterState("csntp", "baseUrl1_")) {
+      ClusterState clusterState = zkStateReader.getClusterState();
+      assertNotNull(clusterState);
+      assertEquals(1, clusterState.getCollectionStates().size());
+      DocCollection collection1 = clusterState.getCollectionOrNull("collection1");
+      assertNotNull(collection1);
+      assertEquals(DocRouter.DEFAULT, collection1.getRouter());
+      assertEquals(1, collection1.getActiveSlices().size());
+      assertEquals(1, collection1.getSlices().size());
+      Slice slice1 = collection1.getSlice("slice1");
+      assertNotNull(slice1);
+      assertEquals(3, slice1.getReplicas().size());
+      assertEquals(1, slice1.getReplicas(replica -> replica.getType() == Replica.Type.NRT).size());
+      assertEquals(1, slice1.getReplicas(replica -> replica.getType() == Replica.Type.TLOG).size());
+      assertEquals(1, slice1.getReplicas(replica -> replica.getType() == Replica.Type.PULL).size());
+    }
+  }
+
+  @Test
+  public void testBuildClusterState_ReplicaStateAndType() {
+    try (ZkStateReader zkStateReader = ClusterStateMockUtil.buildClusterState("csrStRpDnF", "baseUrl1_")) {
+      ClusterState clusterState = zkStateReader.getClusterState();
+      assertNotNull(clusterState);
+      assertEquals(1, clusterState.getCollectionStates().size());
+      DocCollection collection1 = clusterState.getCollectionOrNull("collection1");
+      assertNotNull(collection1);
+      assertEquals(DocRouter.DEFAULT, collection1.getRouter());
+      assertEquals(1, collection1.getActiveSlices().size());
+      assertEquals(1, collection1.getSlices().size());
+      Slice slice1 = collection1.getSlice("slice1");
+      assertNotNull(slice1);
+      assertEquals(4, slice1.getReplicas().size());
+      assertEquals(1, slice1.getReplicas(replica -> replica.getType() == Replica.Type.NRT && replica.getState() == Replica.State.ACTIVE).size());
+      assertEquals(1, slice1.getReplicas(replica -> replica.getType() == Replica.Type.NRT && replica.getState() == Replica.State.RECOVERY_FAILED).size());
+      assertEquals(1, slice1.getReplicas(replica -> replica.getType() == Replica.Type.TLOG && replica.getState() == Replica.State.RECOVERING).size());
+      assertEquals(1, slice1.getReplicas(replica -> replica.getType() == Replica.Type.PULL && replica.getState() == Replica.State.DOWN).size());
+    }
+  }
+}