From 4e5bcc80b1633e9734a8a85fee02bd691734ebbb Mon Sep 17 00:00:00 2001
From: Rishab Nahata <rnnahata@amazon.com>
Date: Wed, 21 Aug 2024 14:57:38 +0530
Subject: [PATCH] Add tests for move

Signed-off-by: Rishab Nahata <rnnahata@amazon.com>
---
 ...TimeBoundBalancedShardsAllocatorTests.java | 40 +++++++++++++++++--
 1 file changed, 37 insertions(+), 3 deletions(-)

diff --git a/server/src/test/java/org/opensearch/cluster/routing/allocation/TimeBoundBalancedShardsAllocatorTests.java b/server/src/test/java/org/opensearch/cluster/routing/allocation/TimeBoundBalancedShardsAllocatorTests.java
index 1c32716d7f138..2f6f030d93f03 100644
--- a/server/src/test/java/org/opensearch/cluster/routing/allocation/TimeBoundBalancedShardsAllocatorTests.java
+++ b/server/src/test/java/org/opensearch/cluster/routing/allocation/TimeBoundBalancedShardsAllocatorTests.java
@@ -206,10 +206,9 @@ public void testAllShardsMoveWhenExcludedAndTimeoutNotBreached() {
         allocator.allocate(allocation);
         List<ShardRouting> relocatingShards = allocation.routingNodes().shardsWithState(ShardRoutingState.RELOCATING);
         assertEquals(node1ShardCount, relocatingShards.size());
-        assertEquals(node1ShardCount, allocation.routingNodes().getRelocatingShardCount());
     }
 
-    public void testNoShardsMoveWhenExcludedAndTimeoutNotBreached() {
+    public void testNoShardsMoveWhenExcludedAndTimeoutBreached() {
         int numberOfIndices = 3;
         int numberOfShards = 5;
         int numberOfReplicas = 1;
@@ -240,7 +239,42 @@ public void testNoShardsMoveWhenExcludedAndTimeoutNotBreached() {
         allocator.allocate(allocation);
         List<ShardRouting> relocatingShards = allocation.routingNodes().shardsWithState(ShardRoutingState.RELOCATING);
         assertEquals(0, relocatingShards.size());
-        assertEquals(0, allocation.routingNodes().getRelocatingShardCount());
+    }
+
+    public void testPartialShardsMoveWhenExcludedAndTimeoutBreached() {
+        int numberOfIndices = 3;
+        int numberOfShards = 5;
+        int numberOfReplicas = 1;
+        int totalShardCount = numberOfIndices * (numberOfShards * (numberOfReplicas + 1));
+        Metadata metadata = buildMetadata(Metadata.builder(), numberOfIndices, numberOfShards, numberOfReplicas);
+        RoutingTable routingTable = buildRoutingTable(metadata);
+        ClusterState state = ClusterState.builder(ClusterName.CLUSTER_NAME_SETTING.getDefault(Settings.EMPTY))
+            .metadata(metadata)
+            .routingTable(routingTable)
+            .nodes(DiscoveryNodes.builder().add(node1).add(node2).add(node3))
+            .build();
+        MockAllocationService allocationService = createAllocationService();
+        state = applyStartedShardsUntilNoChange(state, allocationService);
+        // check all shards allocated
+        assertEquals(0, state.getRoutingNodes().shardsWithState(INITIALIZING).size());
+        assertEquals(totalShardCount, state.getRoutingNodes().shardsWithState(STARTED).size());
+        Settings settings = Settings.builder().put("cluster.routing.allocation.exclude.zone", "1a").build();
+        // since for moves, it creates an iterator over shards which interleaves between nodes, hence
+        // for shardsToMove=6, it will have 2 shards from node1, node2, node3 each attempting to move with only
+        // shards from node1 can actually move. Hence, total moves that will be executed is 2 (6/3).
+        int shardsToMove = 5; // such that time out is never breached
+        BalancedShardsAllocator allocator = new TestBalancedShardsAllocator(settings, new CountDownLatch(shardsToMove));
+        RoutingAllocation allocation = new RoutingAllocation(
+            allocationDecidersForExcludeAPI(settings),
+            new RoutingNodes(state, false),
+            state,
+            ClusterInfo.EMPTY,
+            null,
+            System.nanoTime()
+        );
+        allocator.allocate(allocation);
+        List<ShardRouting> relocatingShards = allocation.routingNodes().shardsWithState(ShardRoutingState.RELOCATING);
+        assertEquals(shardsToMove / 3, relocatingShards.size());
     }
 
     private RoutingTable buildRoutingTable(Metadata metadata) {