From a1b6db26db5e03a31492549a181c285f9b35c9a2 Mon Sep 17 00:00:00 2001
From: Anshum Gupta <anshum@apache.org>
Date: Mon, 10 Sep 2018 14:20:07 -0700
Subject: [PATCH] SOLR-12762: Fix javadoc for SolrCloudTestCase.clusterShape()
 method and add a method that validates only against Active slices while
 testing

---
 solr/CHANGES.txt                              |  6 ++-
 .../apache/solr/cloud/SolrCloudTestCase.java  | 40 ++++++++++++++-----
 2 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/solr/CHANGES.txt b/solr/CHANGES.txt
index 0d71204f5e5..f9e5b562b2f 100644
--- a/solr/CHANGES.txt
+++ b/solr/CHANGES.txt
@@ -84,7 +84,11 @@ Apache ZooKeeper 3.4.11
 Jetty 9.4.11.v20180605
 
 
-(No Changes)
+Other Changes
+----------------------
+
+* SOLR-12762: Fix javadoc for SolrCloudTestCase.clusterShape() method and add a method that validates only against
+  Active slices (Anshum Gupta)
 
 
 ==================  7.5.0 ==================
diff --git a/solr/test-framework/src/java/org/apache/solr/cloud/SolrCloudTestCase.java b/solr/test-framework/src/java/org/apache/solr/cloud/SolrCloudTestCase.java
index 1b5f67bbd51..bd041f00ca7 100644
--- a/solr/test-framework/src/java/org/apache/solr/cloud/SolrCloudTestCase.java
+++ b/solr/test-framework/src/java/org/apache/solr/cloud/SolrCloudTestCase.java
@@ -283,7 +283,7 @@ public class SolrCloudTestCase extends SolrTestCaseJ4 {
 
   /**
    * Return a {@link CollectionStatePredicate} that returns true if a collection has the expected
-   * number of active shards and active replicas
+   * number of shards and active replicas
    */
   public static CollectionStatePredicate clusterShape(int expectedShards, int expectedReplicas) {
     return (liveNodes, collectionState) -> {
@@ -291,19 +291,39 @@ public class SolrCloudTestCase extends SolrTestCaseJ4 {
         return false;
       if (collectionState.getSlices().size() != expectedShards)
         return false;
-      for (Slice slice : collectionState) {
-        int activeReplicas = 0;
-        for (Replica replica : slice) {
-          if (replica.isActive(liveNodes))
-            activeReplicas++;
-        }
-        if (activeReplicas != expectedReplicas)
-          return false;
-      }
+      if (compareActiveReplicaCountsForShards(expectedReplicas, liveNodes, collectionState)) return false;
       return true;
     };
   }
 
+  /**
+   * Return a {@link CollectionStatePredicate} that returns true if a collection has the expected
+   * number of active shards and active replicas
+   */
+  public static CollectionStatePredicate activeClusterShape(int expectedShards, int expectedReplicas) {
+    return (liveNodes, collectionState) -> {
+      if (collectionState == null)
+        return false;
+      if (collectionState.getActiveSlices().size() != expectedShards)
+        return false;
+      if (compareActiveReplicaCountsForShards(expectedReplicas, liveNodes, collectionState)) return false;
+      return true;
+    };
+  }
+
+  private static boolean compareActiveReplicaCountsForShards(int expectedReplicas, Set<String> liveNodes, DocCollection collectionState) {
+    for (Slice slice : collectionState) {
+      int activeReplicas = 0;
+      for (Replica replica : slice) {
+        if (replica.isActive(liveNodes))
+          activeReplicas++;
+      }
+      if (activeReplicas != expectedReplicas)
+        return true;
+    }
+    return false;
+  }
+
   /**
    * Get a (reproducibly) random shard from a {@link DocCollection}
    */