From 9e2206634356e57d5158d80be268d46f9ebb2607 Mon Sep 17 00:00:00 2001
From: Anshu Agarwal <anshukag@amazon.com>
Date: Fri, 7 Oct 2022 10:28:18 +0530
Subject: [PATCH] test to verify the fix

Signed-off-by: Anshu Agarwal <anshukag@amazon.com>
---
 .../cluster/routing/WeightedRoutingIT.java    | 48 +++++++++++++++++++
 1 file changed, 48 insertions(+)

diff --git a/server/src/internalClusterTest/java/org/opensearch/cluster/routing/WeightedRoutingIT.java b/server/src/internalClusterTest/java/org/opensearch/cluster/routing/WeightedRoutingIT.java
index 6cf8292095c6a..55b92aae5754b 100644
--- a/server/src/internalClusterTest/java/org/opensearch/cluster/routing/WeightedRoutingIT.java
+++ b/server/src/internalClusterTest/java/org/opensearch/cluster/routing/WeightedRoutingIT.java
@@ -242,4 +242,52 @@ public void testGetWeightedRouting_WeightsAreSet() throws IOException {
         assertEquals(weightedRouting, weightedRoutingResponse.weights());
         assertEquals("3.0", weightedRoutingResponse.getLocalNodeWeight());
     }
+
+    public void testWeightedRoutingMetadataOnOSProcessRestart() throws IOException {
+        Settings commonSettings = Settings.builder()
+            .put("cluster.routing.allocation.awareness.attributes", "zone")
+            .put("cluster.routing.allocation.awareness.force.zone.values", "a,b,c")
+            .build();
+
+        internalCluster().startNodes(
+            Settings.builder().put(commonSettings).put("node.attr.zone", "a").build(),
+            Settings.builder().put(commonSettings).put("node.attr.zone", "b").build(),
+            Settings.builder().put(commonSettings).put("node.attr.zone", "c").build()
+        );
+
+        logger.info("--> waiting for nodes to form a cluster");
+        ClusterHealthResponse health = client().admin().cluster().prepareHealth().setWaitForNodes("3").execute().actionGet();
+        assertThat(health.isTimedOut(), equalTo(false));
+
+        ensureGreen();
+
+        logger.info("--> setting shard routing weights for weighted round robin");
+        Map<String, Double> weights = Map.of("a", 1.0, "b", 2.0, "c", 3.0);
+        WeightedRouting weightedRouting = new WeightedRouting("zone", weights);
+        // put api call to set weights
+        ClusterPutWeightedRoutingResponse response = client().admin()
+            .cluster()
+            .prepareWeightedRouting()
+            .setWeightedRouting(weightedRouting)
+            .get();
+        assertEquals(response.isAcknowledged(), true);
+
+        ensureStableCluster(3);
+
+        // routing weights are set in cluster metadata
+        assertNotNull(internalCluster().clusterService().state().metadata().weightedRoutingMetadata());
+
+        ensureGreen();
+
+        // Restart a random data node and check that OS process comes healthy
+        try {
+            internalCluster().restartRandomDataNode();
+        } catch (Exception e) {
+            e.printStackTrace();
+        }
+
+        ensureGreen();
+
+        assertNotNull(internalCluster().clusterService().state().metadata().weightedRoutingMetadata());
+    }
 }