From 3b2e7d5fd89636195dce5fa8e06192d3fa9cabda Mon Sep 17 00:00:00 2001
From: Smit Patel <patelsmit32123@gmail.com>
Date: Wed, 22 Jan 2025 17:41:11 +0530
Subject: [PATCH] Added tests

---
 .../allocation/BalanceConfigurationTests.java | 49 +++++++++++++++++++
 .../cluster/OpenSearchAllocationTestCase.java | 15 ++++++
 2 files changed, 64 insertions(+)

diff --git a/server/src/test/java/org/opensearch/cluster/routing/allocation/BalanceConfigurationTests.java b/server/src/test/java/org/opensearch/cluster/routing/allocation/BalanceConfigurationTests.java
index 11cbe89645657..7962496e375f6 100644
--- a/server/src/test/java/org/opensearch/cluster/routing/allocation/BalanceConfigurationTests.java
+++ b/server/src/test/java/org/opensearch/cluster/routing/allocation/BalanceConfigurationTests.java
@@ -250,6 +250,55 @@ public void testPrimaryBalanceWithPreferPrimaryBalanceSetting() {
         assertTrue(balanceFailed <= 1);
     }
 
+    /**
+     * This test verifies primary shard balance is not reset when primary shard balance factors are changed.
+     */
+    public void testPreferPrimaryBalanceSettingRemainsUnchangedWithUpdatingPrimaryBalanceFactors() {
+        final int numberOfNodes = 5;
+        final int numberOfIndices = 5;
+        final int numberOfShards = 25;
+        final int numberOfReplicas = 1;
+        final int numberOfRuns = 5;
+        int balanceFailed = 0;
+
+        Settings settings = getSettingsBuilderForPrimaryBalance().build();
+        ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
+        AllocationService strategy = createAllocationService(settings, clusterSettings,  new TestGatewayAllocator(), SNAPSHOT_INFO_SERVICE_WITH_NO_SHARD_SIZES);
+        for (int i = 0; i < numberOfRuns; i++) {
+            ClusterState clusterState = initCluster(strategy, numberOfIndices, numberOfNodes, numberOfShards, numberOfReplicas);
+            clusterState = removeOneNode(clusterState, strategy);
+            logger.info(ShardAllocations.printShardDistribution(clusterState));
+            try {
+                verifyPerIndexPrimaryBalance(clusterState);
+            } catch (AssertionError e) {
+                balanceFailed++;
+                logger.info("Unexpected assertion failure");
+            }
+        }
+        assertTrue(balanceFailed <= 1);
+
+        // Update settings & apply
+        Settings updatedSettings = getSettingsBuilderForPrimaryBalance()
+            .put(BalancedShardsAllocator.INDEX_BALANCE_FACTOR_SETTING.getKey(), BalancedShardsAllocator.INDEX_BALANCE_FACTOR_SETTING.get(settings) + 0.01f)
+            .put(BalancedShardsAllocator.SHARD_BALANCE_FACTOR_SETTING.getKey(), BalancedShardsAllocator.SHARD_BALANCE_FACTOR_SETTING.get(settings) + 0.01f)
+            .put(BalancedShardsAllocator.PRIMARY_SHARD_REBALANCE_BUFFER.getKey(), BalancedShardsAllocator.PRIMARY_SHARD_REBALANCE_BUFFER.get(settings) + 0.01f)
+            .build();
+        clusterSettings.applySettings(updatedSettings);
+        // Test if primary shard balancing still works
+        for (int i = 0; i < numberOfRuns; i++) {
+            ClusterState clusterState = initCluster(strategy, numberOfIndices, numberOfNodes, numberOfShards, numberOfReplicas);
+            clusterState = removeOneNode(clusterState, strategy);
+            logger.info(ShardAllocations.printShardDistribution(clusterState));
+            try {
+                verifyPerIndexPrimaryBalance(clusterState);
+            } catch (AssertionError e) {
+                balanceFailed++;
+                logger.info("Unexpected assertion failure");
+            }
+        }
+        assertTrue(balanceFailed <= 1);
+    }
+
     /**
      * This test verifies primary shard balance is attained  setting.
      */
diff --git a/test/framework/src/main/java/org/opensearch/cluster/OpenSearchAllocationTestCase.java b/test/framework/src/main/java/org/opensearch/cluster/OpenSearchAllocationTestCase.java
index bc42993ac2096..5fb86f2d73acc 100644
--- a/test/framework/src/main/java/org/opensearch/cluster/OpenSearchAllocationTestCase.java
+++ b/test/framework/src/main/java/org/opensearch/cluster/OpenSearchAllocationTestCase.java
@@ -142,6 +142,21 @@ public static MockAllocationService createAllocationService(
         );
     }
 
+    public static MockAllocationService createAllocationService(
+        Settings settings,
+        ClusterSettings clusterSettings,
+        GatewayAllocator gatewayAllocator,
+        SnapshotsInfoService snapshotsInfoService
+    ) {
+        return new MockAllocationService(
+            randomAllocationDeciders(settings, EMPTY_CLUSTER_SETTINGS, random()),
+            gatewayAllocator,
+            new BalancedShardsAllocator(settings, clusterSettings),
+            EmptyClusterInfoService.INSTANCE,
+            snapshotsInfoService
+        );
+    }
+
     public static AllocationDeciders randomAllocationDeciders(Settings settings, ClusterSettings clusterSettings, Random random) {
         List<AllocationDecider> deciders = new ArrayList<>(
             ClusterModule.createAllocationDeciders(settings, clusterSettings, Collections.emptyList())