diff --git a/server/src/internalClusterTest/java/org/opensearch/indices/settings/SearchOnlyReplicaIT.java b/server/src/internalClusterTest/java/org/opensearch/indices/settings/SearchOnlyReplicaIT.java
index 11eb78d7dd7cf..db34bb113320d 100644
--- a/server/src/internalClusterTest/java/org/opensearch/indices/settings/SearchOnlyReplicaIT.java
+++ b/server/src/internalClusterTest/java/org/opensearch/indices/settings/SearchOnlyReplicaIT.java
@@ -152,10 +152,15 @@ public void testFailoverWithSearchReplicaWhenSearchNodeRestarts() throws Excepti
         String searchNode = internalCluster().startSearchOnlyNode();
         ensureGreen(TEST_INDEX);
         // Restart Search Node
-        internalCluster().restartNode(searchNode);
-        // Ensure search shard is unassigned
-        ensureYellowAndNoInitializingShards(TEST_INDEX);
-        assertActiveSearchShards(0);
+        internalCluster().restartNode(searchNode, new InternalTestCluster.RestartCallback() {
+            @Override
+            public Settings onNodeStopped(String nodeName) throws Exception {
+                // Ensure search shard is unassigned
+                ensureYellowAndNoInitializingShards(TEST_INDEX);
+                assertActiveSearchShards(0);
+                return super.onNodeStopped(nodeName);
+            }
+        });
         // Ensure search shard is recovered
         ensureGreen(TEST_INDEX);
         assertActiveSearchShards(1);