Added tests

opensearch-project · Jan 22, 2025 · 3b2e7d5 · 3b2e7d5
1 parent 95a4058
commit 3b2e7d5
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 0 deletions.
diff --git a/...er/src/test/java/org/opensearch/cluster/routing/allocation/BalanceConfigurationTests.java b/...er/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
@@ -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())