You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@solr.apache.org by br...@apache.org on 2023/11/13 11:24:12 UTC

(solr-sandbox) branch main updated: Improve encryption tests to run on multiple hosts and shards. (#83)

This is an automated email from the ASF dual-hosted git repository.

broustant pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/solr-sandbox.git


The following commit(s) were added to refs/heads/main by this push:
     new 5e89f6e  Improve encryption tests to run on multiple hosts and shards. (#83)
5e89f6e is described below

commit 5e89f6e7b59306af0e2c28868d2ab1f05f263cb9
Author: Bruno Roustant <33...@users.noreply.github.com>
AuthorDate: Mon Nov 13 12:24:06 2023 +0100

    Improve encryption tests to run on multiple hosts and shards. (#83)
---
 .../solr/encryption/EncryptionDirectoryTest.java   |  62 +++++---
 .../solr/encryption/EncryptionHeavyLoadTest.java   |  50 ++++---
 .../encryption/EncryptionRequestHandlerTest.java   | 157 +++++++++++----------
 .../{TestUtil.java => EncryptionTestUtil.java}     |  70 +++++++--
 .../encryption/EncryptionUpdateHandlerTest.java    |   4 +-
 .../solr/encryption/EncryptionUpdateLogTest.java   |   2 +-
 6 files changed, 209 insertions(+), 136 deletions(-)

diff --git a/encryption/src/test/java/org/apache/solr/encryption/EncryptionDirectoryTest.java b/encryption/src/test/java/org/apache/solr/encryption/EncryptionDirectoryTest.java
index 3f965a2..a090e5e 100644
--- a/encryption/src/test/java/org/apache/solr/encryption/EncryptionDirectoryTest.java
+++ b/encryption/src/test/java/org/apache/solr/encryption/EncryptionDirectoryTest.java
@@ -21,6 +21,7 @@ import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexInput;
 import org.apache.solr.client.solrj.impl.CloudSolrClient;
 import org.apache.solr.client.solrj.request.CollectionAdminRequest;
+import org.apache.solr.client.solrj.request.UpdateRequest;
 import org.apache.solr.cloud.MiniSolrCloudCluster;
 import org.apache.solr.cloud.SolrCloudTestCase;
 import org.junit.AfterClass;
@@ -55,15 +56,15 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
 
   private String collectionName;
   private CloudSolrClient solrClient;
-  private TestUtil testUtil;
+  private EncryptionTestUtil testUtil;
 
   @BeforeClass
   public static void beforeClass() throws Exception {
     System.setProperty(PROPERTY_INNER_ENCRYPTION_DIRECTORY_FACTORY, MockFactory.class.getName());
     System.setProperty("solr." + PARAM_KEY_SUPPLIER_FACTORY, TestingKeySupplier.Factory.class.getName());
-    TestUtil.setInstallDirProperty();
-    cluster = new MiniSolrCloudCluster.Builder(1, createTempDir())
-      .addConfig("config", TestUtil.getConfigPath("collection1"))
+    EncryptionTestUtil.setInstallDirProperty();
+    cluster = new MiniSolrCloudCluster.Builder(2, createTempDir())
+      .addConfig("config", EncryptionTestUtil.getConfigPath("collection1"))
       .configure();
   }
 
@@ -79,10 +80,9 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
     super.setUp();
     collectionName = COLLECTION_PREFIX + UUID.randomUUID();
     solrClient = cluster.getSolrClient();
-    solrClient.setDefaultCollection(collectionName);
-    CollectionAdminRequest.createCollection(collectionName, 1, 1).process(solrClient);
-    cluster.waitForActiveCollection(collectionName, 1, 1);
-    testUtil = new TestUtil(solrClient, collectionName);
+    CollectionAdminRequest.createCollection(collectionName, 2, 2).process(solrClient);
+    cluster.waitForActiveCollection(collectionName, 2, 4);
+    testUtil = new EncryptionTestUtil(solrClient, collectionName);
   }
 
   @Override
@@ -112,28 +112,30 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
     // Create 2 index segments without encryption.
     testUtil.indexDocsAndCommit("weather broadcast");
     testUtil.indexDocsAndCommit("sunny weather");
+    testUtil.indexDocsAndCommit("foo");
+    testUtil.indexDocsAndCommit("bar");
     testUtil.assertQueryReturns("weather", 2);
 
     // Verify that without key id, we can reload the index because it is not encrypted.
-    testUtil.reloadCore();
+    testUtil.reloadCores();
     testUtil.assertQueryReturns("weather", 2);
 
     // Set the encryption key id in the commit user data,
     // and run an optimized commit to rewrite the index, now encrypted.
     mockDir.setKeysInCommitUserData(KEY_ID_1);
-    solrClient.optimize();
+    optimizeCommit();
 
     // Verify that without key id, we cannot decrypt the index anymore.
     mockDir.forceClearText = true;
-    testUtil.assertCannotReloadCore();
+    testUtil.assertCannotReloadCores();
     // Verify that with a wrong key id, we cannot decrypt the index.
     mockDir.forceClearText = false;
     mockDir.forceKeySecret = KEY_SECRET_2;
-    testUtil.assertCannotReloadCore();
+    testUtil.assertCannotReloadCores();
     // Verify that with the right key id, we can decrypt the index and search it.
     mockDir.forceKeySecret = null;
     mockDir.expectedKeySecret = KEY_SECRET_1;
-    testUtil.reloadCore();
+    testUtil.reloadCores();
     testUtil.assertQueryReturns("weather", 2);
     testUtil.assertQueryReturns("sunny", 1);
     mockDir.clearMockValues();
@@ -158,18 +160,19 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
     // Create 1 new segment with the same encryption key id.
     mockDir.setKeysInCommitUserData(KEY_ID_1);
     testUtil.indexDocsAndCommit("foggy weather");
+    testUtil.indexDocsAndCommit("boo");
 
     // Verify that without key id, we cannot decrypt the index.
     mockDir.forceClearText = true;
-    testUtil.assertCannotReloadCore();
+    testUtil.assertCannotReloadCores();
     // Verify that with a wrong key id, we cannot decrypt the index.
     mockDir.forceClearText = false;
     mockDir.forceKeySecret = KEY_SECRET_2;
-    testUtil.assertCannotReloadCore();
+    testUtil.assertCannotReloadCores();
     // Verify that with the right key id, we can decrypt the index and search it.
     mockDir.forceKeySecret = null;
     mockDir.expectedKeySecret = KEY_SECRET_1;
-    testUtil.reloadCore();
+    testUtil.reloadCores();
     testUtil.assertQueryReturns("weather", 3);
     testUtil.assertQueryReturns("sunny", 1);
     mockDir.clearMockValues();
@@ -186,19 +189,19 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
     // Set the new encryption key id in the commit user data,
     // and run an optimized commit to rewrite the index, now encrypted with the new key.
     mockDir.setKeysInCommitUserData(KEY_ID_1, KEY_ID_2);
-    solrClient.optimize();
+    optimizeCommit();
 
     // Verify that without key id, we cannot decrypt the index.
     mockDir.forceClearText = true;
-    testUtil.assertCannotReloadCore();
+    testUtil.assertCannotReloadCores();
     // Verify that with a wrong key id, we cannot decrypt the index.
     mockDir.forceClearText = false;
     mockDir.forceKeySecret = KEY_SECRET_1;
-    testUtil.assertCannotReloadCore();
+    testUtil.assertCannotReloadCores();
     // Verify that with the right key id, we can decrypt the index and search it.
     mockDir.forceKeySecret = null;
     mockDir.expectedKeySecret = KEY_SECRET_2;
-    testUtil.reloadCore();
+    testUtil.reloadCores();
     testUtil.assertQueryReturns("weather", 3);
     testUtil.assertQueryReturns("sunny", 1);
   }
@@ -214,15 +217,30 @@ public class EncryptionDirectoryTest extends SolrCloudTestCase {
     // Remove the active key parameter from the commit user data,
     // and run an optimized commit to rewrite the index, now cleartext with no keys.
     mockDir.setKeysInCommitUserData(KEY_ID_1, null);
-    solrClient.optimize();
+    optimizeCommit();
 
     // Verify that without key id, we can reload the index because it is not encrypted.
     mockDir.forceClearText = true;
-    testUtil.reloadCore();
+    testUtil.reloadCores();
     testUtil.assertQueryReturns("weather", 3);
     testUtil.assertQueryReturns("sunny", 1);
   }
 
+  /**
+   * Sends an {@link UpdateRequest} with optimize to all replicas. If there are two segments or more,
+   * then all segments are merged into one, ensuring here that we encrypt all the index data.
+   * <p>
+   * This is not what should be done to encrypt. The real request should be sent to the
+   * {@link EncryptionRequestHandler}, but this test is designed to work independently.
+   */
+  private void optimizeCommit() {
+    testUtil.forAllReplicas(replica -> {
+      UpdateRequest request = new UpdateRequest();
+      request.setAction(UpdateRequest.ACTION.OPTIMIZE, true, true, 1);
+      testUtil.requestCore(request, replica);
+    });
+  }
+
   public static class MockFactory implements EncryptionDirectoryFactory.InnerFactory {
     @Override
     public EncryptionDirectory create(Directory delegate,
diff --git a/encryption/src/test/java/org/apache/solr/encryption/EncryptionHeavyLoadTest.java b/encryption/src/test/java/org/apache/solr/encryption/EncryptionHeavyLoadTest.java
index 5a8608b..aa6be0c 100644
--- a/encryption/src/test/java/org/apache/solr/encryption/EncryptionHeavyLoadTest.java
+++ b/encryption/src/test/java/org/apache/solr/encryption/EncryptionHeavyLoadTest.java
@@ -19,7 +19,6 @@ package org.apache.solr.encryption;
 import com.carrotsearch.randomizedtesting.generators.RandomStrings;
 import org.apache.solr.client.solrj.SolrQuery;
 import org.apache.solr.client.solrj.SolrRequest;
-import org.apache.solr.client.solrj.SolrServerException;
 import org.apache.solr.client.solrj.impl.CloudSolrClient;
 import org.apache.solr.client.solrj.request.CollectionAdminRequest;
 import org.apache.solr.client.solrj.request.GenericSolrRequest;
@@ -35,19 +34,18 @@ import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Random;
 import java.util.Set;
-import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Supplier;
 
 import static org.apache.solr.encryption.EncryptionRequestHandler.*;
+import static org.apache.solr.encryption.EncryptionRequestHandlerTest.EncryptionStatus;
 import static org.apache.solr.encryption.TestingKeySupplier.*;
 
 /**
@@ -78,7 +76,9 @@ public class EncryptionHeavyLoadTest extends SolrCloudTestCase {
   private static final String COLLECTION_PREFIX = EncryptionHeavyLoadTest.class.getSimpleName() + "-collection-";
   private static final String SYSTEM_OUTPUT_MARKER = "*** ";
 
+  private volatile String collectionName;
   private volatile CloudSolrClient solrClient;
+  private volatile EncryptionTestUtil testUtil;
   private volatile boolean stopTest;
   private volatile Dictionary dictionary;
   private List<Thread> threads;
@@ -91,9 +91,9 @@ public class EncryptionHeavyLoadTest extends SolrCloudTestCase {
 
   @BeforeClass
   public static void beforeClass() throws Exception {
-    TestUtil.setInstallDirProperty();
-    cluster = new MiniSolrCloudCluster.Builder(1, createTempDir())
-      .addConfig("config", TestUtil.getConfigPath("collection1"))
+    EncryptionTestUtil.setInstallDirProperty();
+    cluster = new MiniSolrCloudCluster.Builder(2, createTempDir())
+      .addConfig("config", EncryptionTestUtil.getConfigPath("collection1"))
       .configure();
   }
 
@@ -105,11 +105,11 @@ public class EncryptionHeavyLoadTest extends SolrCloudTestCase {
   @Override
   public void setUp() throws Exception {
     super.setUp();
-    String collectionName = COLLECTION_PREFIX + UUID.randomUUID();
+    collectionName = COLLECTION_PREFIX + random().nextLong();
     solrClient = cluster.getSolrClient();
-    solrClient.setDefaultCollection(collectionName);
-    CollectionAdminRequest.createCollection(collectionName, 1, 1).process(solrClient);
-    cluster.waitForActiveCollection(collectionName, 1, 1);
+    CollectionAdminRequest.createCollection(collectionName, 2, 2).process(solrClient);
+    cluster.waitForActiveCollection(collectionName, 2, 4);
+    testUtil = new EncryptionTestUtil(solrClient, collectionName);
     dictionary = new Dictionary.Builder().build(DICTIONARY_SIZE, random());
     threads = new ArrayList<>();
   }
@@ -191,18 +191,18 @@ public class EncryptionHeavyLoadTest extends SolrCloudTestCase {
   }
 
   private boolean encrypt(String keyId, boolean waitForCompletion) throws Exception {
-    NamedList<Object> response = sendEncryptionRequest(keyId);
-    if (response.get(ENCRYPTION_STATE).equals(STATE_PENDING)) {
+    EncryptionStatus encryptionStatus = sendEncryptionRequests(keyId);
+    if (!encryptionStatus.complete) {
       if (!waitForCompletion) {
         return false;
       }
       print("waiting for encryption completion for keyId=" + keyId);
-      while (response.get(ENCRYPTION_STATE).equals(STATE_PENDING)) {
+      while (!encryptionStatus.complete) {
         if (isTimeElapsed()) {
           return false;
         }
         Thread.sleep(500);
-        response = sendEncryptionRequest(keyId);
+        encryptionStatus = sendEncryptionRequests(keyId);
       }
       print("encryption complete for keyId=" + keyId);
     }
@@ -213,12 +213,18 @@ public class EncryptionHeavyLoadTest extends SolrCloudTestCase {
     return random.nextFloat() <= PROBABILITY_OF_WAITING_ENCRYPTION_COMPLETION;
   }
 
-  private NamedList<Object> sendEncryptionRequest(String keyId) throws SolrServerException, IOException {
+  private EncryptionStatus sendEncryptionRequests(String keyId) {
     ModifiableSolrParams params = new ModifiableSolrParams();
     params.set(PARAM_KEY_ID, keyId);
-    NamedList<Object> response = solrClient.request(new GenericSolrRequest(SolrRequest.METHOD.GET, "/admin/encrypt", params));
-    print("encrypt keyId=" + keyId + " => response status=" + response.get(STATUS) + " state=" + response.get(ENCRYPTION_STATE));
-    return response;
+    GenericSolrRequest encryptRequest = new GenericSolrRequest(SolrRequest.METHOD.GET, "/admin/encrypt", params);
+    EncryptionStatus encryptionStatus = new EncryptionStatus();
+    testUtil.forAllReplicas(replica -> {
+      NamedList<Object> response = testUtil.requestCore(encryptRequest, replica);
+      encryptionStatus.success &= response.get(STATUS).equals(STATUS_SUCCESS);
+      encryptionStatus.complete &= response.get(ENCRYPTION_STATE).equals(STATE_COMPLETE);
+    });
+    print("encrypt keyId=" + keyId + " => response success=" + encryptionStatus.success + " complete=" + encryptionStatus.complete);
+    return encryptionStatus;
   }
 
   private static void print(String message) {
@@ -246,7 +252,7 @@ public class EncryptionHeavyLoadTest extends SolrCloudTestCase {
       Dictionary build(int size, Random random) {
         Set<String> terms = new HashSet<>();
         for (int i = 0; i < size;) {
-          String term = RandomStrings.randomUnicodeOfCodepointLengthBetween(random, 4, 12);
+          String term = RandomStrings.randomAsciiLettersOfLengthBetween(random, 4, 12);
           if (terms.add(term)) {
             i++;
           }
@@ -279,10 +285,10 @@ public class EncryptionHeavyLoadTest extends SolrCloudTestCase {
             docs.add(createDoc(random));
           }
           totalDocs += docs.size();
-          solrClient.add(docs);
+          solrClient.add(collectionName, docs);
           if (random.nextFloat() <= PROBABILITY_OF_COMMIT_PER_BATCH) {
             numCommits++;
-            solrClient.commit();
+            solrClient.commit(collectionName);
           }
           if (++numBatches % 10 == 0) {
             threadPrint("sent " + numBatches + " indexing batches, totalDocs=" + totalDocs + ", numCommits=" + numCommits);
@@ -330,7 +336,7 @@ public class EncryptionHeavyLoadTest extends SolrCloudTestCase {
           QueryResponse response = null;
           do {
             try {
-              response = solrClient.query(new SolrQuery(dictionary.getTerm(random)));
+              response = solrClient.query(collectionName, new SolrQuery(dictionary.getTerm(random)));
             } catch (Exception e) {
               // Some queries might not be parseable due to the random terms. Just retry with another term.
             }
diff --git a/encryption/src/test/java/org/apache/solr/encryption/EncryptionRequestHandlerTest.java b/encryption/src/test/java/org/apache/solr/encryption/EncryptionRequestHandlerTest.java
index 515d8b7..ee8f087 100644
--- a/encryption/src/test/java/org/apache/solr/encryption/EncryptionRequestHandlerTest.java
+++ b/encryption/src/test/java/org/apache/solr/encryption/EncryptionRequestHandlerTest.java
@@ -52,18 +52,19 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
 
   private static final String COLLECTION_PREFIX = EncryptionRequestHandlerTest.class.getSimpleName() + "-collection-";
 
-  private static MockEncryptionDirectory mockDir;
+  private static volatile boolean forceClearText;
+  private static volatile String soleKeyIdAllowed;
 
   private String collectionName;
   private CloudSolrClient solrClient;
-  private TestUtil testUtil;
+  private EncryptionTestUtil testUtil;
 
   @BeforeClass
   public static void beforeClass() throws Exception {
     System.setProperty(PROPERTY_INNER_ENCRYPTION_DIRECTORY_FACTORY, MockFactory.class.getName());
-    TestUtil.setInstallDirProperty();
-    cluster = new MiniSolrCloudCluster.Builder(1, createTempDir())
-      .addConfig("config", TestUtil.getConfigPath("collection1"))
+    EncryptionTestUtil.setInstallDirProperty();
+    cluster = new MiniSolrCloudCluster.Builder(2, createTempDir())
+      .addConfig("config", EncryptionTestUtil.getConfigPath("collection1"))
       .configure();
   }
 
@@ -78,15 +79,14 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
     super.setUp();
     collectionName = COLLECTION_PREFIX + UUID.randomUUID();
     solrClient = cluster.getSolrClient();
-    solrClient.setDefaultCollection(collectionName);
-    CollectionAdminRequest.createCollection(collectionName, 1, 1).process(solrClient);
-    cluster.waitForActiveCollection(collectionName, 1, 1);
-    testUtil = new TestUtil(solrClient, collectionName);
+    CollectionAdminRequest.createCollection(collectionName, 2, 2).process(solrClient);
+    cluster.waitForActiveCollection(collectionName, 2, 4);
+    testUtil = new EncryptionTestUtil(solrClient, collectionName);
   }
 
   @Override
   public void tearDown() throws Exception {
-    mockDir.clearMockValues();
+    clearMockValues();
     CollectionAdminRequest.deleteCollection(collectionName).process(solrClient);
     super.tearDown();
   }
@@ -94,39 +94,39 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
   @Test
   public void testEncryptionFromNoKeysToOneKey_NoIndex() throws Exception {
     // Send an encrypt request with a key id on an empty index.
-    NamedList<Object> response = encrypt(KEY_ID_1);
-    assertEquals(STATUS_SUCCESS, response.get(STATUS));
-    assertEquals(STATE_COMPLETE, response.get(ENCRYPTION_STATE));
+    EncryptionStatus encryptionStatus = encrypt(KEY_ID_1);
+    assertTrue(encryptionStatus.success);
+    assertTrue(encryptionStatus.complete);
 
     // Index some documents to create a first segment.
     testUtil.indexDocsAndCommit("weather broadcast");
 
     // Verify that the segment is encrypted.
-    mockDir.forceClearText = true;
-    testUtil.assertCannotReloadCore();
-    mockDir.forceClearText = false;
-    testUtil.reloadCore();
+    forceClearText = true;
+    testUtil.assertCannotReloadCores();
+    forceClearText = false;
+    testUtil.reloadCores();
     testUtil.assertQueryReturns("weather", 1);
   }
 
   @Test
   public void testEncryptionFromNoKeysToOneKeyToNoKeys_NoIndex() throws Exception {
     // Send an encrypt request with a key id on an empty index.
-    NamedList<Object> response = encrypt(KEY_ID_1);
-    assertEquals(STATUS_SUCCESS, response.get(STATUS));
-    assertEquals(STATE_COMPLETE, response.get(ENCRYPTION_STATE));
+    EncryptionStatus encryptionStatus = encrypt(KEY_ID_1);
+    assertTrue(encryptionStatus.success);
+    assertTrue(encryptionStatus.complete);
 
     // Send another encrypt request with no key id, still on the empty index.
-    response = encrypt(NO_KEY_ID);
-    assertEquals(STATUS_SUCCESS, response.get(STATUS));
-    assertEquals(STATE_COMPLETE, response.get(ENCRYPTION_STATE));
+    encryptionStatus = encrypt(NO_KEY_ID);
+    assertTrue(encryptionStatus.success);
+    assertTrue(encryptionStatus.complete);
 
     // Index some documents to create a first segment.
     testUtil.indexDocsAndCommit("weather broadcast");
 
     // Verify that the segment is cleartext.
-    mockDir.forceClearText = true;
-    testUtil.reloadCore();
+    forceClearText = true;
+    testUtil.reloadCores();
     testUtil.assertQueryReturns("weather", 1);
   }
 
@@ -140,26 +140,26 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
     testUtil.indexDocsAndCommit("weather broadcast");
     testUtil.indexDocsAndCommit("sunny weather");
     // Verify that the segments are cleartext.
-    mockDir.forceClearText = true;
-    testUtil.reloadCore();
+    forceClearText = true;
+    testUtil.reloadCores();
     testUtil.assertQueryReturns("weather", 2);
-    mockDir.forceClearText = false;
+    forceClearText = false;
 
     // Send an encrypt request with a key id.
-    NamedList<Object> response = encrypt(KEY_ID_1);
-    assertEquals(STATUS_SUCCESS, response.get(STATUS));
-    assertEquals(STATE_PENDING, response.get(ENCRYPTION_STATE));
+    EncryptionStatus encryptionStatus = encrypt(KEY_ID_1);
+    assertTrue(encryptionStatus.success);
+    assertFalse(encryptionStatus.complete);
 
     waitUntilEncryptionIsComplete(KEY_ID_1);
 
     // Verify that the segment is encrypted.
-    mockDir.forceClearText = true;
-    testUtil.assertCannotReloadCore();
-    mockDir.forceClearText = false;
-    mockDir.soleKeyIdAllowed = KEY_ID_1;
-    testUtil.reloadCore();
+    forceClearText = true;
+    testUtil.assertCannotReloadCores();
+    forceClearText = false;
+    soleKeyIdAllowed = KEY_ID_1;
+    testUtil.reloadCores();
     testUtil.assertQueryReturns("weather", 2);
-    mockDir.clearMockValues();
+    clearMockValues();
   }
 
   @Test
@@ -170,18 +170,18 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
     testUtil.indexDocsAndCommit("foggy weather");
 
     // Send an encrypt request with another key id.
-    NamedList<Object> response = encrypt(KEY_ID_2);
-    assertEquals(STATUS_SUCCESS, response.get(STATUS));
-    assertEquals(STATE_PENDING, response.get(ENCRYPTION_STATE));
+    EncryptionStatus encryptionStatus = encrypt(KEY_ID_2);
+    assertTrue(encryptionStatus.success);
+    assertFalse(encryptionStatus.complete);
 
     waitUntilEncryptionIsComplete(KEY_ID_2);
 
     // Verify that the segment is encrypted.
-    mockDir.forceClearText = true;
-    testUtil.assertCannotReloadCore();
-    mockDir.forceClearText = false;
-    mockDir.soleKeyIdAllowed = KEY_ID_2;
-    testUtil.reloadCore();
+    forceClearText = true;
+    testUtil.assertCannotReloadCores();
+    forceClearText = false;
+    soleKeyIdAllowed = KEY_ID_2;
+    testUtil.reloadCores();
     testUtil.assertQueryReturns("weather", 3);
   }
 
@@ -193,41 +193,48 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
     testUtil.indexDocsAndCommit("foggy weather");
 
     // Send an encrypt request with no key id.
-    NamedList<Object> response = encrypt(NO_KEY_ID);
-    assertEquals(STATUS_SUCCESS, response.get(STATUS));
-    assertEquals(STATE_PENDING, response.get(ENCRYPTION_STATE));
+    EncryptionStatus encryptionStatus = encrypt(NO_KEY_ID);
+    assertTrue(encryptionStatus.success);
+    assertFalse(encryptionStatus.complete);
 
     waitUntilEncryptionIsComplete(NO_KEY_ID);
 
     // Verify that the segment is cleartext.
-    mockDir.forceClearText = true;
-    testUtil.reloadCore();
+    forceClearText = true;
+    testUtil.reloadCores();
     testUtil.assertQueryReturns("weather", 3);
-    mockDir.clearMockValues();
+    clearMockValues();
 
     // Index some documents to ensure we have at least two segments.
     testUtil.indexDocsAndCommit("cloudy weather");
 
     // Send an encrypt request with another key id.
-    response = encrypt(KEY_ID_2);
-    assertEquals(STATUS_SUCCESS, response.get(STATUS));
-    assertEquals(STATE_PENDING, response.get(ENCRYPTION_STATE));
+    encryptionStatus = encrypt(KEY_ID_2);
+    assertTrue(encryptionStatus.success);
+    assertFalse(encryptionStatus.complete);
 
     waitUntilEncryptionIsComplete(KEY_ID_2);
 
     // Verify that the segment is encrypted.
-    mockDir.forceClearText = true;
-    testUtil.assertCannotReloadCore();
-    mockDir.forceClearText = false;
-    mockDir.soleKeyIdAllowed = KEY_ID_2;
-    testUtil.reloadCore();
+    forceClearText = true;
+    testUtil.assertCannotReloadCores();
+    forceClearText = false;
+    soleKeyIdAllowed = KEY_ID_2;
+    testUtil.reloadCores();
     testUtil.assertQueryReturns("weather", 4);
   }
 
-  private NamedList<Object> encrypt(String keyId) throws Exception {
+  private EncryptionStatus encrypt(String keyId) {
     ModifiableSolrParams params = new ModifiableSolrParams();
     params.set(PARAM_KEY_ID, keyId);
-    return solrClient.request(new GenericSolrRequest(SolrRequest.METHOD.GET, "/admin/encrypt", params));
+    GenericSolrRequest encryptRequest = new GenericSolrRequest(SolrRequest.METHOD.GET, "/admin/encrypt", params);
+    EncryptionStatus encryptionStatus = new EncryptionStatus();
+    testUtil.forAllReplicas(replica -> {
+      NamedList<Object> response = testUtil.requestCore(encryptRequest, replica);
+      encryptionStatus.success &= response.get(STATUS).equals(STATUS_SUCCESS);
+      encryptionStatus.complete &= response.get(ENCRYPTION_STATE).equals(STATE_COMPLETE);
+    });
+    return encryptionStatus;
   }
 
   private void waitUntilEncryptionIsComplete(String keyId) throws InterruptedException {
@@ -236,41 +243,38 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
                          100,
                          TimeUnit.MILLISECONDS,
                          () -> {
-                           NamedList<Object> response;
+                           EncryptionStatus encryptionStatus;
                            try {
-                             response = encrypt(keyId);
+                             encryptionStatus = encrypt(keyId);
                            } catch (Exception e) {
                              throw new RuntimeException(e);
                            }
-                           assertEquals(STATUS_SUCCESS, response.get(STATUS));
-                           return response.get(ENCRYPTION_STATE).equals(STATE_COMPLETE);
+                           assertTrue(encryptionStatus.success);
+                           return encryptionStatus.complete;
                          });
   }
 
+  private static void clearMockValues() {
+    forceClearText = false;
+    soleKeyIdAllowed = null;
+  }
+
   public static class MockFactory implements EncryptionDirectoryFactory.InnerFactory {
     @Override
     public EncryptionDirectory create(Directory delegate,
                                       AesCtrEncrypterFactory encrypterFactory,
                                       KeySupplier keySupplier) throws IOException {
-      return mockDir = new MockEncryptionDirectory(delegate, encrypterFactory, keySupplier);
+      return new MockEncryptionDirectory(delegate, encrypterFactory, keySupplier);
     }
   }
 
   private static class MockEncryptionDirectory extends EncryptionDirectory {
 
-    boolean forceClearText;
-    String soleKeyIdAllowed;
-
     MockEncryptionDirectory(Directory delegate, AesCtrEncrypterFactory encrypterFactory, KeySupplier keySupplier)
       throws IOException {
       super(delegate, encrypterFactory, keySupplier);
     }
 
-    void clearMockValues() {
-      forceClearText = false;
-      soleKeyIdAllowed = null;
-    }
-
     @Override
     public IndexInput openInput(String fileName, IOContext context) throws IOException {
       return forceClearText ? in.openInput(fileName, context) : super.openInput(fileName, context);
@@ -301,4 +305,9 @@ public class EncryptionRequestHandlerTest extends SolrCloudTestCase {
       }
     }
   }
+
+  public static class EncryptionStatus {
+    public boolean success = true;
+    public boolean complete = true;
+  }
 }
diff --git a/encryption/src/test/java/org/apache/solr/encryption/TestUtil.java b/encryption/src/test/java/org/apache/solr/encryption/EncryptionTestUtil.java
similarity index 56%
rename from encryption/src/test/java/org/apache/solr/encryption/TestUtil.java
rename to encryption/src/test/java/org/apache/solr/encryption/EncryptionTestUtil.java
index 9dba03c..d5c9128 100644
--- a/encryption/src/test/java/org/apache/solr/encryption/TestUtil.java
+++ b/encryption/src/test/java/org/apache/solr/encryption/EncryptionTestUtil.java
@@ -18,28 +18,35 @@ package org.apache.solr.encryption;
 
 import org.apache.solr.SolrTestCaseJ4;
 import org.apache.solr.client.solrj.SolrQuery;
+import org.apache.solr.client.solrj.SolrRequest;
+import org.apache.solr.client.solrj.SolrServerException;
 import org.apache.solr.client.solrj.impl.CloudSolrClient;
+import org.apache.solr.client.solrj.impl.Http2SolrClient;
 import org.apache.solr.client.solrj.request.CoreAdminRequest;
 import org.apache.solr.client.solrj.response.QueryResponse;
 import org.apache.solr.common.SolrException;
 import org.apache.solr.common.SolrInputDocument;
-import org.apache.solr.common.cloud.DocCollection;
+import org.apache.solr.common.cloud.Replica;
 import org.apache.solr.common.cloud.Slice;
+import org.apache.solr.common.params.CoreAdminParams;
+import org.apache.solr.common.util.NamedList;
 import org.junit.Assert;
 
+import java.io.IOException;
 import java.nio.file.Path;
+import java.util.function.Consumer;
 
 /**
  * Utility methods for encryption tests.
  */
-public class TestUtil {
+public class EncryptionTestUtil {
 
-  private final CloudSolrClient solrClient;
+  private final CloudSolrClient cloudSolrClient;
   private final String collectionName;
   private int docId;
 
-  public TestUtil(CloudSolrClient solrClient, String collectionName) {
-    this.solrClient = solrClient;
+  public EncryptionTestUtil(CloudSolrClient cloudSolrClient, String collectionName) {
+    this.cloudSolrClient = cloudSolrClient;
     this.collectionName = collectionName;
   }
 
@@ -72,44 +79,77 @@ public class TestUtil {
       SolrInputDocument doc = new SolrInputDocument();
       doc.addField("id", Integer.toString(docId++));
       doc.addField("text", text);
-      solrClient.add(doc);
+      cloudSolrClient.add(collectionName, doc);
     }
-    solrClient.commit(collectionName);
+    cloudSolrClient.commit(collectionName);
   }
 
   /**
    * Verifies that the provided query returns the expected number of results.
    */
   public void assertQueryReturns(String query, int expectedNumResults) throws Exception {
-    QueryResponse response = solrClient.query(new SolrQuery(query));
+    QueryResponse response = cloudSolrClient.query(collectionName, new SolrQuery(query));
     Assert.assertEquals(expectedNumResults, response.getResults().size());
   }
 
   /**
    * Reloads the leader replica core of the first shard of the collection.
    */
-  public void reloadCore() throws Exception {
+  public void reloadCores() throws Exception {
     try {
-      DocCollection collection = solrClient.getClusterState().getCollection(collectionName);
-      Slice slice = collection.getSlices().iterator().next();
-      CoreAdminRequest.reloadCore(slice.getLeader().core, solrClient);
+      forAllReplicas(replica -> {
+        try {
+          CoreAdminRequest req = new CoreAdminRequest();
+          req.setBasePath(replica.getBaseUrl());
+          req.setCoreName(replica.getCoreName());
+          req.setAction(CoreAdminParams.CoreAdminAction.RELOAD);
+          try (Http2SolrClient httpSolrClient = new Http2SolrClient.Builder(replica.getBaseUrl()).build()) {
+            httpSolrClient.request(req);
+          }
+        } catch (SolrServerException e) {
+          throw new SolrException(SolrException.ErrorCode.SERVER_ERROR, e);
+        } catch (IOException e) {
+          throw new RuntimeException(e);
+        }
+      });
     } catch (SolrException e) {
       throw new CoreReloadException("The index cannot be reloaded. There is probably an issue with the encryption key ids.", e);
     }
   }
 
   /**
-   * Verifies that {@link #reloadCore()} fails.
+   * Verifies that {@link #reloadCores()} fails.
    */
-  public void assertCannotReloadCore() throws Exception {
+  public void assertCannotReloadCores() throws Exception {
     try {
-      reloadCore();
+      reloadCores();
       Assert.fail("Core reloaded whereas it was not expected to be possible");
     } catch (CoreReloadException e) {
       // Expected.
     }
   }
 
+  /** Processes the given {@code action} for all replicas of the collection defined in the constructor. */
+  public void forAllReplicas(Consumer<Replica> action) {
+    for (Slice slice : cloudSolrClient.getClusterState().getCollection(collectionName).getSlices()) {
+      for (Replica replica : slice.getReplicas()) {
+        action.accept(replica);
+      }
+    }
+  }
+
+  /** Sends the given {@link SolrRequest} to a specific replica. */
+  public NamedList<Object> requestCore(SolrRequest<?> request, Replica replica) {
+    request.setBasePath(replica.getCoreUrl());
+    try (Http2SolrClient httpSolrClient = new Http2SolrClient.Builder(replica.getBaseUrl()).build()) {
+      return httpSolrClient.request(request);
+    } catch (SolrServerException e) {
+      throw new SolrException(SolrException.ErrorCode.SERVER_ERROR, e);
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
   private static class CoreReloadException extends Exception {
     CoreReloadException(String msg, SolrException cause) {
       super(msg, cause);
diff --git a/encryption/src/test/java/org/apache/solr/encryption/EncryptionUpdateHandlerTest.java b/encryption/src/test/java/org/apache/solr/encryption/EncryptionUpdateHandlerTest.java
index cf87cd5..c5720f6 100644
--- a/encryption/src/test/java/org/apache/solr/encryption/EncryptionUpdateHandlerTest.java
+++ b/encryption/src/test/java/org/apache/solr/encryption/EncryptionUpdateHandlerTest.java
@@ -37,8 +37,8 @@ public class EncryptionUpdateHandlerTest extends SolrTestCaseJ4 {
 
   @BeforeClass
   public static void beforeClass() throws Exception {
-    TestUtil.setInstallDirProperty();
-    initCore("solrconfig.xml", "schema.xml", TestUtil.getConfigPath().toString());
+    EncryptionTestUtil.setInstallDirProperty();
+    initCore("solrconfig.xml", "schema.xml", EncryptionTestUtil.getConfigPath().toString());
   }
 
   /**
diff --git a/encryption/src/test/java/org/apache/solr/encryption/EncryptionUpdateLogTest.java b/encryption/src/test/java/org/apache/solr/encryption/EncryptionUpdateLogTest.java
index 9757830..1dec37c 100644
--- a/encryption/src/test/java/org/apache/solr/encryption/EncryptionUpdateLogTest.java
+++ b/encryption/src/test/java/org/apache/solr/encryption/EncryptionUpdateLogTest.java
@@ -63,7 +63,7 @@ public class EncryptionUpdateLogTest extends SolrCloudTestCase {
   @BeforeClass
   public static void setupClass() throws Exception {
     configureCluster(NUM_SHARDS * NUM_REPLICAS)
-      .addConfig("config", TestUtil.getConfigPath("collection1"))
+      .addConfig("config", EncryptionTestUtil.getConfigPath("collection1"))
       .configure();
   }