[PLAT-15483] Fix Replica count logic to work with older k8s cluster w…

…here availableReplicas are not present

Summary: [PLAT-15483] Fix Replica count logic to work with older k8s cluster where availableReplicas are not present

Test Plan: Manual testing

Reviewers: sneelakantan

Reviewed By: sneelakantan

Subscribers: sneelakantan, yugaware

Differential Revision: https://phorge.dev.yugabyte.com/D38447

managed/src/main/java/com/yugabyte/yw/common/KubernetesManager.java

@@ -864,6 +864,9 @@ public abstract String getStorageClass(
 
   public abstract String getKubeconfigCluster(Map<String, String> config);
 
+  public abstract boolean checkStatefulSetStatus(
+      Map<String, String> config, String namespace, String labelSelector, int replicaCount);
+
   public abstract void deleteUnusedPVCs(
       Map<String, String> config,
       String namespace,

managed/src/main/java/com/yugabyte/yw/common/NativeKubernetesManager.java

@@ -599,4 +599,10 @@ public List<Service> getNamespacedServices(
     // TODO Auto-generated method stub
     throw new UnsupportedOperationException("Unimplemented method 'getNamespacedServices'");
   }
+
+  @Override
+  public boolean checkStatefulSetStatus(
+      Map<String, String> config, String namespace, String labelSelector, int replicaCount) {
+    throw new UnsupportedOperationException("Not implemented");
+  }
 }

managed/src/main/java/com/yugabyte/yw/common/ShellKubernetesManager.java

@@ -516,7 +516,8 @@ public List<PersistentVolumeClaim> getPVCs(
     return deserialize(response.getMessage(), PersistentVolumeClaimList.class).getItems();
   }
 
-  private boolean checkStatefulSetStatus(
+  @Override
+  public boolean checkStatefulSetStatus(
       Map<String, String> config, String namespace, String labelSelector, int replicaCount) {
     List<String> commandList =
         ImmutableList.of(
@@ -544,22 +545,51 @@ private boolean checkStatefulSetStatus(
             "get",
             "statefulset",
             statefulSetNames[0],
-            "-o=jsonpath={.status.availableReplicas} {.status.replicas}");
+            "-o=jsonpath=replicas={.status.replicas}|readyReplicas={.status.readyReplicas}|availableReplicas={.status.availableReplicas}");
     response =
         execCommand(config, commandList, false)
             .processErrors("Unable to get StatefulSet status for " + statefulSetNames[0]);
 
-    // 2 values in output
-    String[] replicaCounts = response.getMessage().trim().split(" ");
-    boolean isReady = false;
-    if (replicaCounts.length == 2) {
-      int availableReplicas = Integer.parseInt(replicaCounts[0]);
-      int totalReplicas = Integer.parseInt(replicaCounts[1]);
-      if (availableReplicas == totalReplicas && totalReplicas == replicaCount) {
-        isReady = true;
+    Map<String, Integer> parsedValues = parseKubectlOutput(response.getMessage());
+
+    // Access values from the map
+    int replicas = parsedValues.get("replicas");
+    int readyReplicas = parsedValues.get("readyReplicas");
+    int availableReplicas = parsedValues.get("availableReplicas");
+
+    if (replicas <= 0 || replicas == availableReplicas || replicas == readyReplicas) {
+      // Either no replicas or all replicas are available/ready
+      return true;
+    }
+    return false;
+  }
+
+  private static Map<String, Integer> parseKubectlOutput(String response) {
+    // Create a map to store the parsed key-value pairs as integers
+    Map<String, Integer> resultMap = new HashMap<>();
+
+    resultMap.put("replicas", -1);
+    resultMap.put("readyReplicas", -1);
+    resultMap.put("availableReplicas", -1);
+
+    String[] fields = response.split("\\|");
+
+    for (String field : fields) {
+      // Split each field by '=' to separate key and value
+      String[] keyValue = field.split("=", 2);
+
+      // Check if we have both key and value
+      if (keyValue.length == 2 && !keyValue[1].isEmpty()) {
+        try {
+          // Parse the value as an integer and store it in the map
+          resultMap.put(keyValue[0], Integer.parseInt(keyValue[1]));
+        } catch (NumberFormatException e) {
+          // If parsing fails, keep the default value of -1
+          resultMap.put(keyValue[0], -1);
+        }
       }
     }
-    return isReady;
+    return resultMap;
   }
 
   @Override

managed/src/test/java/com/yugabyte/yw/common/KubernetesManagerTest.java

@@ -37,6 +37,7 @@
 import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.Mockito;
+import org.mockito.Mockito.*;
 import org.mockito.MockitoAnnotations;
 
 @RunWith(JUnitParamsRunner.class)
@@ -241,4 +242,69 @@ public void testGetStatefulSetServerTypeGflagsChecksum(String outputFilePath) th
       }
     }
   }
+
+  @Test
+  public void testCheckStatefulSetStatus_Failure_ReplicaMismatch() {
+    ShellResponse response1 = ShellResponse.create(0, "statefulset1");
+    ShellResponse response2 =
+        ShellResponse.create(0, "replicas=3|readyReplicas=1|availableReplicas=1");
+    Map<String, String> testConfig = new HashMap<String, String>();
+    when(shellProcessHandler.run(anyList(), any(ShellProcessContext.class))).thenReturn(response1);
+    when(shellProcessHandler.run(anyList(), any(ShellProcessContext.class))).thenReturn(response2);
+    boolean status =
+        kubernetesManager.checkStatefulSetStatus(testConfig, "test-ns", "test-release", 3);
+    assertEquals(false, status);
+  }
+
+  @Test
+  public void testCheckStatefulSetStatus_success() {
+    ShellResponse response1 = ShellResponse.create(0, "statefulset1");
+    ShellResponse response2 =
+        ShellResponse.create(0, "replicas=3|readyReplicas=3|availableReplicas=3");
+    Map<String, String> testConfig = new HashMap<String, String>();
+    when(shellProcessHandler.run(anyList(), any(ShellProcessContext.class))).thenReturn(response1);
+    when(shellProcessHandler.run(anyList(), any(ShellProcessContext.class))).thenReturn(response2);
+    boolean status =
+        kubernetesManager.checkStatefulSetStatus(testConfig, "test-ns", "test-release", 3);
+    assertEquals(true, status);
+  }
+
+  @Test
+  public void testCheckStatefulSetStatus_success_malformed_1() {
+    ShellResponse response1 = ShellResponse.create(0, "statefulset1");
+    ShellResponse response2 =
+        ShellResponse.create(0, "replicas=3|readyReplicas=3|availableReplicas=");
+    Map<String, String> testConfig = new HashMap<String, String>();
+    when(shellProcessHandler.run(anyList(), any(ShellProcessContext.class))).thenReturn(response1);
+    when(shellProcessHandler.run(anyList(), any(ShellProcessContext.class))).thenReturn(response2);
+    boolean status =
+        kubernetesManager.checkStatefulSetStatus(testConfig, "test-ns", "test-release", 3);
+    assertEquals(true, status);
+  }
+
+  @Test
+  public void testCheckStatefulSetStatus_success_malformed_2() {
+    ShellResponse response1 = ShellResponse.create(0, "statefulset1");
+    ShellResponse response2 =
+        ShellResponse.create(0, "replicas=3|readyReplicas=|availableReplicas=3");
+    Map<String, String> testConfig = new HashMap<String, String>();
+    when(shellProcessHandler.run(anyList(), any(ShellProcessContext.class))).thenReturn(response1);
+    when(shellProcessHandler.run(anyList(), any(ShellProcessContext.class))).thenReturn(response2);
+    boolean status =
+        kubernetesManager.checkStatefulSetStatus(testConfig, "test-ns", "test-release", 3);
+    assertEquals(true, status);
+  }
+
+  @Test
+  public void testCheckStatefulSetStatus_success_no_replicas() {
+    ShellResponse response1 = ShellResponse.create(0, "statefulset1");
+    ShellResponse response2 =
+        ShellResponse.create(0, "replicas=0|readyReplicas=|availableReplicas=3");
+    Map<String, String> testConfig = new HashMap<String, String>();
+    when(shellProcessHandler.run(anyList(), any(ShellProcessContext.class))).thenReturn(response1);
+    when(shellProcessHandler.run(anyList(), any(ShellProcessContext.class))).thenReturn(response2);
+    boolean status =
+        kubernetesManager.checkStatefulSetStatus(testConfig, "test-ns", "test-release", 3);
+    assertEquals(true, status);
+  }
 }