From 4c07b0b7c36455171a27a121817473e6b6c1029b Mon Sep 17 00:00:00 2001
From: Akram Ben Aissi <akram.benaissi@gmail.com>
Date: Mon, 16 Dec 2024 16:39:12 +0100
Subject: [PATCH] Refactoring apply

---
 src/codeflare_sdk/ray/cluster/cluster.py | 39 ++++++++++--------------
 1 file changed, 16 insertions(+), 23 deletions(-)

diff --git a/src/codeflare_sdk/ray/cluster/cluster.py b/src/codeflare_sdk/ray/cluster/cluster.py
index e2b989ca6..c6074afda 100644
--- a/src/codeflare_sdk/ray/cluster/cluster.py
+++ b/src/codeflare_sdk/ray/cluster/cluster.py
@@ -194,40 +194,33 @@ def apply(self, force=False):
         # check if RayCluster CustomResourceDefinition exists if not throw RuntimeError
         self._throw_for_no_raycluster()
         namespace = self.config.namespace
-
+        name = self.config.name
         try:
             self.config_check()
             api_instance = client.CustomObjectsApi(get_api_client())
             crds = self.get_dynamic_client().resources
-            api_instance = crds.get(
-                api_version="workload.codeflare.dev/v1beta2", kind="AppWrapper"
-            )
             if self.config.appwrapper:
+                api_version = "workload.codeflare.dev/v1beta2"
+                api_instance = crds.get(api_version, kind="AppWrapper")
+                # defaulting body to resource_yaml
+                body = self.resource_yaml
                 if self.config.write_to_file:
+                    # if write_to_file is True, load the file from AppWrapper yaml and update body
                     with open(self.resource_yaml) as f:
                         aw = yaml.load(f, Loader=yaml.FullLoader)
-                        api_instance.server_side_apply(
-                            group="workload.codeflare.dev",
-                            version="v1beta2",
-                            namespace=namespace,
-                            plural="appwrappers",
-                            body=aw,
-                        )
-                else:
-                    api_instance.server_side_apply(
-                        group="workload.codeflare.dev",
-                        version="v1beta2",
-                        namespace=namespace,
-                        plural="appwrappers",
-                        body=self.resource_yaml,
-                    )
-                print(f"AppWrapper: '{self.config.name}' has successfully been created")
+                    body = aw
+                api_instance.server_side_apply(
+                    group="workload.codeflare.dev",
+                    version="v1beta2",
+                    namespace=namespace,
+                    plural="appwrappers",
+                    body=body,
+                )
+                print(f"AppWrapper: '{name}' has successfully been created")
             else:
                 api_instance = crds.get(api_version="ray.io/v1", kind="RayCluster")
                 self._component_resources_apply(namespace, api_instance)
-                print(
-                    f"Ray Cluster: '{self.config.name}' has successfully been applied"
-                )
+                print(f"Ray Cluster: '{name}' has successfully been applied")
         except Exception as e:  # pragma: no cover
             return _kube_api_error_handling(e)