Refactor node drain function, fix compact cluster drain race

tests/chaos/oadp/conftest.py

@@ -19,7 +19,7 @@
 from utilities.infra import ExecCommandOnPod, wait_for_node_status
 from utilities.oadp import VeleroBackup, create_rhel_vm
 from utilities.storage import write_file
-from utilities.virt import node_mgmt_console, wait_for_node_schedulable_status
+from utilities.virt import drain_node, wait_for_node_schedulable_status
 
 LOGGER = logging.getLogger(__name__)
 
@@ -79,9 +79,13 @@ def rebooted_vm_source_node(rhel_vm_with_dv_running, oadp_backup_in_progress, wo
 
 
 @pytest.fixture()
-def drain_vm_source_node(admin_client, rhel_vm_with_dv_running, oadp_backup_in_progress):
-    vm_node = rhel_vm_with_dv_running.vmi.node
-    with node_mgmt_console(admin_client=admin_client, node=vm_node, node_mgmt="drain"):
+def drain_vm_source_node(
+    admin_client, hco_namespace, compact_cluster, rhel_vm_with_dv_running, oadp_backup_in_progress
+):
+    vm_node = rhel_vm_with_dv_running.vmi.get_node(privileged_client=admin_client)
+    with drain_node(
+        admin_client=admin_client, node=vm_node, hco_namespace=hco_namespace, compact_cluster=compact_cluster
+    ):
         wait_for_node_schedulable_status(node=vm_node, status=False)
         yield vm_node
 

tests/virt/cluster/migration_and_maintenance/test_evictionstrategy.py

@@ -15,7 +15,7 @@
 from utilities.hco import ResourceEditorValidateHCOReconcile
 from utilities.virt import (
     check_migration_process_after_node_drain,
-    node_mgmt_console,
+    drain_node,
     restart_vm_wait_for_running_vm,
     wait_for_node_schedulable_status,
 )
@@ -46,9 +46,11 @@ def assert_vm_restarts_after_node_drain(source_node, vmi, vmi_old_uid):
 
 
 @pytest.fixture()
-def drained_node(admin_client, vm_for_test_from_template_scope_class):
+def drained_node(admin_client, hco_namespace, compact_cluster, vm_for_test_from_template_scope_class):
     source_node = vm_for_test_from_template_scope_class.vmi.get_node(privileged_client=admin_client)
-    with node_mgmt_console(admin_client=admin_client, node=source_node, node_mgmt="drain"):
+    with drain_node(
+        admin_client=admin_client, node=source_node, hco_namespace=hco_namespace, compact_cluster=compact_cluster
+    ):
         yield source_node
 
 

tests/virt/node/migration_and_maintenance/test_node_maintenance.py

@@ -21,9 +21,10 @@
 from utilities.virt import (
     VirtualMachineForTests,
     check_migration_process_after_node_drain,
+    cordon_node,
+    drain_node,
     fedora_vm_body,
     fetch_pid_from_windows_vm,
-    node_mgmt_console,
     running_vm,
     start_and_fetch_processid_on_windows_vm,
 )
@@ -38,19 +39,23 @@
 LOGGER = logging.getLogger(__name__)
 
 
-def drain_using_console(admin_client, source_node, vm):
+def drain_using_console(admin_client, hco_namespace, compact_cluster, source_node, vm):
     with running_sleep_in_linux(vm=vm):
-        with node_mgmt_console(admin_client=admin_client, node=source_node, node_mgmt="drain"):
+        with drain_node(
+            admin_client=admin_client, node=source_node, hco_namespace=hco_namespace, compact_cluster=compact_cluster
+        ):
             check_migration_process_after_node_drain(client=admin_client, vm=vm, admin_client=admin_client)
 
 
-def drain_using_console_windows(admin_client, source_node, vm):
+def drain_using_console_windows(admin_client, hco_namespace, compact_cluster, source_node, vm):
     process_name = OS_PROC_NAME["windows"]
     pre_migrate_processid = start_and_fetch_processid_on_windows_vm(
         vm=vm,
         process_name=process_name,
     )
-    with node_mgmt_console(admin_client=admin_client, node=source_node, node_mgmt="drain"):
+    with drain_node(
+        admin_client=admin_client, node=source_node, hco_namespace=hco_namespace, compact_cluster=compact_cluster
+    ):
         check_migration_process_after_node_drain(client=admin_client, vm=vm, admin_client=admin_client)
         post_migrate_processid = fetch_pid_from_windows_vm(vm=vm, process_name=process_name)
         assert post_migrate_processid == pre_migrate_processid, (
@@ -118,10 +123,18 @@ def migration_job_sampler(client, namespace):
 @pytest.mark.polarion("CNV-3006")
 def test_node_drain_using_console_fedora(
     admin_client,
+    hco_namespace,
+    compact_cluster,
     vm_container_disk_fedora,
 ):
     source_node = vm_container_disk_fedora.vmi.get_node(privileged_client=admin_client)
-    drain_using_console(admin_client=admin_client, source_node=source_node, vm=vm_container_disk_fedora)
+    drain_using_console(
+        admin_client=admin_client,
+        hco_namespace=hco_namespace,
+        compact_cluster=compact_cluster,
+        source_node=source_node,
+        vm=vm_container_disk_fedora,
+    )
 
 
 @pytest.mark.parametrize(
@@ -144,17 +157,25 @@ class TestNodeMaintenanceRHEL:
     def test_node_drain_using_console_rhel(
         self,
         admin_client,
+        hco_namespace,
+        compact_cluster,
         vm_for_test_from_template_scope_class,
     ):
         vm = vm_for_test_from_template_scope_class
         drain_using_console(
-            admin_client=admin_client, source_node=vm.vmi.get_node(privileged_client=admin_client), vm=vm
+            admin_client=admin_client,
+            hco_namespace=hco_namespace,
+            compact_cluster=compact_cluster,
+            source_node=vm.vmi.get_node(privileged_client=admin_client),
+            vm=vm,
         )
 
     @pytest.mark.polarion("CNV-4995")
     def test_migration_when_multiple_nodes_unschedulable_using_console_rhel(
         self,
         admin_client,
+        hco_namespace,
+        compact_cluster,
         schedulable_nodes,
         vm_for_test_from_template_scope_class,
     ):
@@ -177,9 +198,13 @@ def test_migration_when_multiple_nodes_unschedulable_using_console_rhel(
             pod=virt_launcher_pod,
             schedulable_nodes=schedulable_nodes,
         )
-        with node_mgmt_console(admin_client=admin_client, node=cordon_nodes[0], node_mgmt="cordon"):
+        with cordon_node(admin_client=admin_client, node=cordon_nodes[0]):
             drain_using_console(
-                admin_client=admin_client, source_node=vm.vmi.get_node(privileged_client=admin_client), vm=vm
+                admin_client=admin_client,
+                hco_namespace=hco_namespace,
+                compact_cluster=compact_cluster,
+                source_node=vm.vmi.get_node(privileged_client=admin_client),
+                vm=vm,
             )
 
 
@@ -204,11 +229,17 @@ class TestNodeCordonAndDrain:
     def test_node_drain_template_windows(
         self,
         admin_client,
+        hco_namespace,
+        compact_cluster,
         vm_for_test_from_template_scope_class,
     ):
         vm = vm_for_test_from_template_scope_class
         drain_using_console_windows(
-            admin_client=admin_client, source_node=vm.vmi.get_node(privileged_client=admin_client), vm=vm
+            admin_client=admin_client,
+            hco_namespace=hco_namespace,
+            compact_cluster=compact_cluster,
+            source_node=vm.vmi.get_node(privileged_client=admin_client),
+            vm=vm,
         )
 
     @pytest.mark.polarion("CNV-4906")
@@ -218,11 +249,7 @@ def test_node_cordon_template_windows(
         vm_for_test_from_template_scope_class,
     ):
         vm = vm_for_test_from_template_scope_class
-        with node_mgmt_console(
-            admin_client=admin_client,
-            node=vm.vmi.get_node(privileged_client=admin_client),
-            node_mgmt="cordon",
-        ):
+        with cordon_node(admin_client=admin_client, node=vm.vmi.get_node(privileged_client=admin_client)):
             with pytest.raises(TimeoutExpiredError):
                 migration_job_sampler(
                     client=admin_client,

tests/virt/node/migration_and_maintenance/test_post_copy_migration.py

@@ -20,10 +20,10 @@
 )
 from utilities.virt import (
     check_migration_process_after_node_drain,
+    drain_node,
     fetch_pid_from_linux_vm,
     fetch_pid_from_windows_vm,
     migrate_vm_and_verify,
-    node_mgmt_console,
     start_and_fetch_processid_on_linux_vm,
     start_and_fetch_processid_on_windows_vm,
 )
@@ -78,9 +78,12 @@ def migrated_hotplugged_vm(hotplugged_vm):
 
 
 @pytest.fixture()
-def drained_node_with_hotplugged_vm(admin_client, hotplugged_vm):
-    with node_mgmt_console(
-        admin_client=admin_client, node=hotplugged_vm.vmi.get_node(privileged_client=admin_client), node_mgmt="drain"
+def drained_node_with_hotplugged_vm(admin_client, hco_namespace, compact_cluster, hotplugged_vm):
+    with drain_node(
+        admin_client=admin_client,
+        node=hotplugged_vm.vmi.get_node(privileged_client=admin_client),
+        hco_namespace=hco_namespace,
+        compact_cluster=compact_cluster,
     ):
         check_migration_process_after_node_drain(client=admin_client, vm=hotplugged_vm, admin_client=admin_client)
     clean_up_migration_jobs(client=admin_client, vm=hotplugged_vm)

tests/virt/node/migration_and_maintenance/test_vm_unscheduled_node.py

@@ -4,7 +4,7 @@
 from tests.os_params import RHEL_LATEST, RHEL_LATEST_LABELS
 from tests.virt.utils import build_node_affinity_dict
 from utilities.virt import (
-    node_mgmt_console,
+    cordon_node,
     vm_instance_from_template,
     wait_for_node_schedulable_status,
 )
@@ -57,7 +57,7 @@ def test_schedule_vm_on_cordoned_node(admin_client, worker_node1, unscheduled_no
     7. Verify that the VMI is running on the expected node (worker_node1).
     """
 
-    with node_mgmt_console(admin_client=admin_client, node=worker_node1, node_mgmt="cordon"):
+    with cordon_node(admin_client=admin_client, node=worker_node1):
         wait_for_node_schedulable_status(node=worker_node1, status=False)
         unscheduled_node_vm.start()
     unscheduled_node_vm.vmi.wait_for_status(status=VirtualMachineInstance.Status.RUNNING)

utilities/virt.py

@@ -9,11 +9,11 @@
 import secrets
 import shlex
 from collections import defaultdict
+from collections.abc import Generator
 from contextlib import contextmanager
 from copy import deepcopy
 from functools import cache
 from json import JSONDecodeError
-from subprocess import run
 from typing import TYPE_CHECKING, Any
 
 import jinja2
@@ -87,6 +87,7 @@
     TIMEOUT_12MIN,
     TIMEOUT_25MIN,
     TIMEOUT_30MIN,
+    VIRT_API,
     VIRT_HANDLER,
     VIRT_LAUNCHER,
     VIRTCTL,
@@ -2145,27 +2146,80 @@ def vm_instance_from_template(
         yield vm
 
 
+def _uncordon_and_stabilize(admin_client: DynamicClient, node: Node, hco_namespace: str) -> None:
+    """
+    Uncordon a node and wait for KubeVirt to stabilize.
+
+    Args:
+        admin_client: Admin Kubernetes client
+        node: Node to uncordon
+        hco_namespace: HCO namespace
+    """
+    LOGGER.info(f"Uncordon node {node.name}")
+    run_command(command=shlex.split(f"oc adm uncordon {node.name}"))
+    wait_for_node_schedulable_status(node=node, status=True)
+    wait_for_kv_stabilize(admin_client=admin_client, hco_namespace=hco_namespace)
+
+
 @contextmanager
-def node_mgmt_console(admin_client, node, node_mgmt):
+def cordon_node(admin_client: DynamicClient, node: Node) -> Generator[None]:
+    """
+    Cordon a node and uncordon it on exit.
+
+    Args:
+        admin_client: Admin Kubernetes client
+        node: Node to cordon
+
+    Yields:
+        None: Control returns while node is cordoned, uncordon happens on exit.
+    """
     hco_namespace = get_hco_namespace(admin_client=admin_client)
     try:
-        LOGGER.info(f"{node_mgmt.capitalize()} the node {node.name}")
-        extra_opts = "--delete-emptydir-data --ignore-daemonsets=true --force" if node_mgmt == "drain" else ""
-        run(
-            f"nohup oc adm {node_mgmt} {node.name} {extra_opts} &",
-            shell=True,
-        )
+        LOGGER.info(f"Cordon the node {node.name}")
+        run_command(command=shlex.split(f"oc adm cordon {node.name}"))
         yield
     finally:
-        if node_mgmt == "drain":
-            LOGGER.info("Terminate drain process")
-            run(
-                shlex.split('pkill -f "oc adm drain"'),
-            )
-        LOGGER.info(f"Uncordon node {node.name}")
-        run(f"oc adm uncordon {node.name}", shell=True)
-        wait_for_node_schedulable_status(node=node, status=True)
-        wait_for_kv_stabilize(admin_client=admin_client, hco_namespace=hco_namespace)
+        _uncordon_and_stabilize(admin_client=admin_client, node=node, hco_namespace=hco_namespace)
+
+
+@contextmanager
+def drain_node(
+    admin_client: DynamicClient, node: Node, hco_namespace: str, compact_cluster: bool = False
+) -> Generator[None]:
+    """
+    Drain a node and uncordon it on exit.
+
+    On compact clusters, relocates virt-api pods before drain to avoid webhook race conditions.
+
+    Args:
+        admin_client: Admin Kubernetes client
+        node: Node to drain
+        hco_namespace: HCO namespace
+        compact_cluster: If True, relocate virt-api pods before drain.
+    """
+    if compact_cluster:
+        for pod in utilities.infra.get_pods(
+            client=admin_client,
+            namespace=hco_namespace,
+            label=f"{Pod.ApiGroup.KUBEVIRT_IO}={VIRT_API}",
+        ):
+            if pod.node.name == node.name:
+                LOGGER.info(
+                    f"Compact cluster: cordoning {node.name} and deleting virt-api pod {pod.name} "
+                    "before drain to avoid webhook race"
+                )
+                with cordon_node(admin_client=admin_client, node=node):
+                    pod.delete(wait=True)
+
+    try:
+        LOGGER.info(f"Drain the node {node.name}")
+        cmd = f"nohup oc adm drain {node.name} --delete-emptydir-data --ignore-daemonsets=true --force &>/dev/null &"
+        run_command(command=["/bin/bash", "-c", cmd])
+        yield
+    finally:
+        LOGGER.info("Terminate drain process")
+        run_command(command=shlex.split('pkill -f "oc adm drain"'), check=False, verify_stderr=False)
+        _uncordon_and_stabilize(admin_client=admin_client, node=node, hco_namespace=hco_namespace)
 
 
 @contextmanager