updating to use krkn-lib virt functions

Assisted By: Claude Code

Signed-off-by: Paige Patton <prubenda@redhat.com>

Author: paigerube14

krkn/scenario_plugins/kubevirt_vm_outage/kubevirt_vm_outage_scenario_plugin.py

@@ -46,7 +46,7 @@ def run(
         try:
             with open(scenario, "r") as f:
                 scenario_config = yaml.full_load(f)
-                
+            
             self.init_clients(lib_telemetry.get_lib_kubernetes())
             pods_status = PodsStatus()
             for config in scenario_config["scenarios"]:
@@ -71,67 +71,6 @@ def init_clients(self, k8s_client: KrknKubernetes):
         self.custom_object_client = k8s_client.custom_object_client
         logging.info("Successfully initialized Kubernetes client for KubeVirt operations")
 
-    def get_vmi(self, name: str, namespace: str) -> Optional[Dict]:
-        """
-        Get a Virtual Machine Instance by name and namespace.
-        
-        :param name: Name of the VMI to retrieve
-        :param namespace: Namespace of the VMI
-        :return: The VMI object if found, None otherwise
-        """
-        try:
-            vmi = self.custom_object_client.get_namespaced_custom_object(
-                group="kubevirt.io",
-                version="v1",
-                namespace=namespace,
-                plural="virtualmachineinstances",
-                name=name
-            )
-            return vmi
-        except ApiException as e:
-            if e.status == 404:
-                logging.warning(f"VMI {name} not found in namespace {namespace}")
-                return None
-            else:
-                logging.error(f"Error getting VMI {name}: {e}")
-                raise
-        except Exception as e:
-            logging.error(f"Unexpected error getting VMI {name}: {e}")
-            raise
-            
-    def get_vmis(self, regex_name: str, namespace: str) -> Optional[Dict]:
-        """
-        Get a Virtual Machine Instance by name and namespace.
-        
-        :param name: Name of the VMI to retrieve
-        :param namespace: Namespace of the VMI
-        :return: The VMI object if found, None otherwise
-        """
-        try:
-            namespaces = self.k8s_client.list_namespaces_by_regex(namespace)
-            for namespace in namespaces:
-                vmis = self.custom_object_client.list_namespaced_custom_object(
-                    group="kubevirt.io",
-                    version="v1",
-                    namespace=namespace,
-                    plural="virtualmachineinstances",
-                )
-
-                for vmi in vmis.get("items"):
-                    vmi_name = vmi.get("metadata",{}).get("name")
-                    match = re.match(regex_name, vmi_name)
-                    if match:
-                        self.vmis_list.append(vmi)
-        except ApiException as e:
-            if e.status == 404:
-                logging.warning(f"VMI {regex_name} not found in namespace {namespace}")
-                return []
-            else:
-                logging.error(f"Error getting VMI {regex_name}: {e}")
-                raise
-        except Exception as e:
-            logging.error(f"Unexpected error getting VMI {regex_name}: {e}")
-            raise
 
     def execute_scenario(self, config: Dict[str, Any], scenario_telemetry: ScenarioTelemetry) -> int:
         """
@@ -154,7 +93,7 @@ def execute_scenario(self, config: Dict[str, Any], scenario_telemetry: ScenarioT
                 logging.error("vm_name parameter is required")
                 return 1
             self.pods_status = PodsStatus()
-            self.get_vmis(vm_name,namespace)
+            self.vmis_list = self.k8s_client.get_vmis(vm_name,namespace)
             for _ in range(kill_count):
 
                 rand_int = random.randint(0, len(self.vmis_list) - 1)
@@ -166,7 +105,7 @@ def execute_scenario(self, config: Dict[str, Any], scenario_telemetry: ScenarioT
                 if not self.validate_environment(vmi_name, vmi_namespace):
                     return 1
 
-                vmi = self.get_vmi(vmi_name, vmi_namespace)
+                vmi = self.k8s_client.get_vmi(vmi_name, vmi_namespace)
                 self.affected_pod = AffectedPod(
                     pod_name=vmi_name,
                     namespace=vmi_namespace,
@@ -212,15 +151,13 @@ def validate_environment(self, vm_name: str, namespace: str) -> bool:
         """
         try:
             # Check if KubeVirt CRDs exist
-            crd_list = self.custom_object_client.list_namespaced_custom_object("kubevirt.io","v1",namespace,"virtualmachines")
-            kubevirt_crds = [crd for crd in crd_list.items() ]
-            
+            kubevirt_crds = self.k8s_client.get_vms(vm_name, namespace)
             if not kubevirt_crds:
                 logging.error("KubeVirt CRDs not found. Ensure KubeVirt/CNV is installed in the cluster")
                 return False
 
             # Check if VMI exists
-            vmi = self.get_vmi(vm_name, namespace)
+            vmi = self.k8s_client.get_vmi(vm_name, namespace)
             if not vmi:
                 logging.error(f"VMI {vm_name} not found in namespace {namespace}")
                 return False
@@ -243,28 +180,15 @@ def patch_vm_spec(self, vm_name: str, namespace: str, running: bool) -> bool:
         """
         try:
             # Get the VM object first to get its current spec
-            vm = self.custom_object_client.get_namespaced_custom_object(
-                group="kubevirt.io",
-                version="v1",
-                namespace=namespace,
-                plural="virtualmachines",
-                name=vm_name
-            )
+            vm = self.k8s_client.get_vm(vm_name, namespace)
 
             # Update the running state
             if 'spec' not in vm:
                 vm['spec'] = {}
             vm['spec']['running'] = running
 
             # Apply the patch
-            self.custom_object_client.patch_namespaced_custom_object(
-                group="kubevirt.io",
-                version="v1",
-                namespace=namespace,
-                plural="virtualmachines",
-                name=vm_name,
-                body=vm
-            )
+            self.k8s_client.patch_vm(vm_name,namespace,vm)
             return True
 
         except ApiException as e:
@@ -293,26 +217,12 @@ def delete_vmi(self, vm_name: str, namespace: str, disable_auto_restart: bool =
                                " - proceeding with deletion but VM may auto-restart")
             start_creation_time =  self.original_vmi.get('metadata', {}).get('creationTimestamp')
             start_time = time.time()
-            try:
-                self.custom_object_client.delete_namespaced_custom_object(
-                    group="kubevirt.io",
-                    version="v1",
-                    namespace=namespace,
-                    plural="virtualmachineinstances",
-                    name=vm_name
-                )
-            except ApiException as e:
-                if e.status == 404:
-                    logging.warning(f"VMI {vm_name} not found during deletion")
-                    return 1
-                else:
-                    logging.error(f"API error during VMI deletion: {e}")
-                    return 1
+            self.k8s_client.delete_vmi(vm_name, namespace)
 
             # Wait for the VMI to be deleted
 
             while time.time() - start_time < timeout:
-                deleted_vmi = self.get_vmi(vm_name, namespace)
+                deleted_vmi = self.k8s_client.get_vmi(vm_name, namespace)
                 if deleted_vmi:
                     if start_creation_time != deleted_vmi.get('metadata', {}).get('creationTimestamp'):
                         logging.info(f"VMI {vm_name} successfully recreated")
@@ -337,7 +247,7 @@ def wait_for_running(self, vm_name: str, namespace: str, timeout: int = 120) ->
         while time.time() - start_time < timeout: 
 
             # Check current state once since we've already waited for the duration
-            vmi = self.get_vmi(vm_name, namespace)
+            vmi = self.k8s_client.get_vmi(vm_name, namespace)
 
             if vmi:
                 if vmi.get('status', {}).get('phase') == "Running":
@@ -378,13 +288,7 @@ def recover(self, vm_name: str, namespace: str, disable_auto_restart: bool = Fal
                                 del metadata[field]
 
                     # Create the VMI
-                    self.custom_object_client.create_namespaced_custom_object(
-                        group="kubevirt.io",
-                        version="v1",
-                        namespace=namespace,
-                        plural="virtualmachineinstances",
-                        body=vmi_dict
-                    )
+                    self.k8s_client.create_vmi(vm_name, namespace, vmi_dict)
                     logging.info(f"Successfully recreated VMI {vm_name}")
 
                     # Wait for VMI to start running

krkn/utils/VirtChecker.py

@@ -40,13 +40,13 @@ def __init__(self, kubevirt_check_config, iterations, krkn_lib: KrknKubernetes,
             self.kube_vm_plugin = KubevirtVmOutageScenarioPlugin()
             self.kube_vm_plugin.init_clients(k8s_client=krkn_lib)
 
-            self.kube_vm_plugin.get_vmis(vmi_name_match,self.namespace)
+            self.vmis_list = self.kube_vm_plugin.k8s_client.get_vmis(vmi_name_match,self.namespace)
         except Exception as e:
             logging.error('Virt Check init exception: ' + str(e))
             return
         # See if multiple node names exist
         node_name_list = [node_name for node_name in self.node_names.split(',') if node_name]
-        for vmi in self.kube_vm_plugin.vmis_list:
+        for vmi in self.vmis_list:
             node_name = vmi.get("status",{}).get("nodeName")
             vmi_name = vmi.get("metadata",{}).get("name")
             ip_address = vmi.get("status",{}).get("interfaces",[])[0].get("ipAddress")