From 7a479ea4a96891f38413cdabcca88bc5b445faa9 Mon Sep 17 00:00:00 2001
From: Chin-Ya Huang <chin-ya.huang@suse.com>
Date: Thu, 5 Dec 2024 17:47:56 +0800
Subject: [PATCH] refactor(robot): generalize expand workload volume keyword

longhorn/longhorn-6633

Signed-off-by: Chin-Ya Huang <chin-ya.huang@suse.com>
---
 e2e/keywords/statefulset.resource      |  4 ----
 e2e/keywords/workload.resource         |  6 ++++++
 e2e/libs/keywords/workload_keywords.py |  8 +++++---
 e2e/libs/utility/utility.py            | 15 +++++++++++++++
 4 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/e2e/keywords/statefulset.resource b/e2e/keywords/statefulset.resource
index 6ebf77bda8..2fd0e8f2af 100644
--- a/e2e/keywords/statefulset.resource
+++ b/e2e/keywords/statefulset.resource
@@ -27,10 +27,6 @@ Scale up statefulset ${statefulset_id} to attach volume
     ${statefulset_name} =    generate_name_with_suffix    statefulset    ${statefulset_id}
     scale_statefulset_up    ${statefulset_name}
 
-Expand statefulset ${statefulset_id} volume by ${size} MiB
-    ${statefulset_name} =    generate_name_with_suffix    statefulset    ${statefulset_id}
-    expand_workload_claim_size_by_mib    ${statefulset_name}    ${size}
-
 Write ${size} MB data to file ${file_name} in statefulset ${statefulset_id}
     ${statefulset_name} =    generate_name_with_suffix    statefulset    ${statefulset_id}
     write_workload_pod_random_data    ${statefulset_name}    ${size}    ${file_name}
diff --git a/e2e/keywords/workload.resource b/e2e/keywords/workload.resource
index a5b90b98d6..c7d64daddf 100644
--- a/e2e/keywords/workload.resource
+++ b/e2e/keywords/workload.resource
@@ -213,3 +213,9 @@ Delete Longhorn ${workload_kind} ${workload_name} pod
     ${pod_name} =    get_workload_pod_name    ${workload_name}    longhorn-system
     Log    ${pod_name}
     delete_pod    ${pod_name}     longhorn-system
+
+Expand ${workload_kind} ${workload_id} volume by ${size}
+    ${workload_name} =   generate_name_with_suffix    ${workload_kind}    ${workload_id}
+    ${new_size} =    convert_size_to_bytes    ${size}
+
+    expand_workload_claim_size    ${workload_name}    ${new_size}
diff --git a/e2e/libs/keywords/workload_keywords.py b/e2e/libs/keywords/workload_keywords.py
index d34fb1567d..a5e23c25b6 100644
--- a/e2e/libs/keywords/workload_keywords.py
+++ b/e2e/libs/keywords/workload_keywords.py
@@ -28,6 +28,7 @@
 from utility.constant import ANNOT_CHECKSUM
 from utility.constant import ANNOT_EXPANDED_SIZE
 from utility.constant import LABEL_LONGHORN_COMPONENT
+from utility.utility import convert_size_to_bytes
 from utility.utility import logging
 
 from volume import Volume
@@ -158,11 +159,12 @@ def wait_for_workload_volume_detached(self, workload_name):
         logging(f'Waiting for {workload_name} volume {volume_name} to be detached')
         self.volume.wait_for_volume_detached(volume_name)
 
-    def expand_workload_claim_size_by_mib(self, workload_name, size_in_mib, claim_index=0):
+    def expand_workload_claim_size(self, workload_name, size_in_byte, claim_index=0):
         claim_name = get_workload_persistent_volume_claim_name(workload_name, index=claim_index)
-        size_in_byte = int(size_in_mib) * MEBIBYTE
+        current_size = self.persistentvolumeclaim.get(claim_name).spec.resources.requests['storage']
+        current_size_byte = convert_size_to_bytes(current_size)
 
-        logging(f'Expanding {workload_name} persistentvolumeclaim {claim_name} by {size_in_mib} MiB')
+        logging(f'Expanding {workload_name} persistentvolumeclaim {claim_name} from {current_size_byte} to {size_in_byte}')
         self.persistentvolumeclaim.expand(claim_name, size_in_byte)
 
     def wait_for_workload_claim_size_expanded(self, workload_name, claim_index=0):
diff --git a/e2e/libs/utility/utility.py b/e2e/libs/utility/utility.py
index bfe8073d77..d7f49d99cd 100644
--- a/e2e/libs/utility/utility.py
+++ b/e2e/libs/utility/utility.py
@@ -338,3 +338,18 @@ def get_name_suffix(*args):
         if arg:
             suffix += f"-{arg}"
     return suffix
+
+
+def convert_size_to_bytes(size):
+    size = size.replace(" ", "")
+
+    if size.endswith("GiB"):
+        return int(size[:-3]) * 1024 * 1024 * 1024
+
+    if size.endswith("MiB"):
+        return int(size[:-3]) * 1024 * 1024
+
+    if size.isdigit():
+        return int(size)
+
+    assert False, f"Invalid size format: {size}"