Refactor: Export SetNodeUnschedulable() and SetNodeSchedulable() func…

…tions

Signed-off-by: Itamar Holder <[email protected]>

tests/libnode/BUILD.bazel

@@ -6,6 +6,7 @@ go_library(
     importpath = "kubevirt.io/kubevirt/tests/libnode",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/util/nodes:go_default_library",
         "//pkg/util/types:go_default_library",
         "//pkg/virt-config:go_default_library",
         "//pkg/virt-controller/services:go_default_library",

tests/libnode/node.go

@@ -26,6 +26,8 @@ import (
 	"strings"
 	"time"
 
+	"kubevirt.io/kubevirt/pkg/util/nodes"
+
 	utiltype "kubevirt.io/kubevirt/pkg/util/types"
 
 	"k8s.io/apimachinery/pkg/api/equality"
@@ -333,3 +335,29 @@ func GetArch() string {
 	Expect(nodes).ToNot(BeEmpty(), "There should be some node")
 	return nodes[0].Status.NodeInfo.Architecture
 }
+
+func setNodeSchedualability(nodeName string, virtCli kubecli.KubevirtClient, setSchedulable bool) {
+	origNode, err := virtCli.CoreV1().Nodes().Get(context.Background(), nodeName, k8smetav1.GetOptions{})
+	Expect(err).ShouldNot(HaveOccurred())
+
+	nodeCopy := origNode.DeepCopy()
+	nodeCopy.Spec.Unschedulable = !setSchedulable
+
+	err = nodes.PatchNode(virtCli, origNode, nodeCopy)
+	Expect(err).ShouldNot(HaveOccurred())
+
+	Eventually(func() bool {
+		patchedNode, err := virtCli.CoreV1().Nodes().Get(context.Background(), nodeName, k8smetav1.GetOptions{})
+		Expect(err).ShouldNot(HaveOccurred())
+
+		return patchedNode.Spec.Unschedulable
+	}, 30*time.Second, time.Second).Should(Equal(!setSchedulable), fmt.Sprintf("node %s is expected to set to Unschedulable=%t, but it's set to %t", nodeName, !setSchedulable, setSchedulable))
+}
+
+func SetNodeUnschedulable(nodeName string, virtCli kubecli.KubevirtClient) {
+	setNodeSchedualability(nodeName, virtCli, false)
+}
+
+func SetNodeSchedulable(nodeName string, virtCli kubecli.KubevirtClient) {
+	setNodeSchedualability(nodeName, virtCli, true)
+}

tests/virt_control_plane_test.go

@@ -129,45 +129,17 @@ var _ = Describe("[Serial][ref_id:2717][sig-compute]KubeVirt control plane resil
 			).Should(BeTrue())
 		}
 
-		setNodeUnschedulable := func(nodeName string) {
-			Eventually(func() error {
-				selectedNode, err := virtCli.CoreV1().Nodes().Get(context.Background(), nodeName, metav1.GetOptions{})
-				if err != nil {
-					return err
-				}
-				selectedNode.Spec.Unschedulable = true
-				if _, err = virtCli.CoreV1().Nodes().Update(context.Background(), selectedNode, metav1.UpdateOptions{}); err != nil {
-					return err
-				}
-				return nil
-			}, 30*time.Second, time.Second).ShouldNot(HaveOccurred())
-		}
-
-		setNodeSchedulable := func(nodeName string) {
-			Eventually(func() error {
-				selectedNode, err := virtCli.CoreV1().Nodes().Get(context.Background(), nodeName, metav1.GetOptions{})
-				if err != nil {
-					return err
-				}
-				selectedNode.Spec.Unschedulable = false
-				if _, err = virtCli.CoreV1().Nodes().Update(context.Background(), selectedNode, metav1.UpdateOptions{}); err != nil {
-					return err
-				}
-				return nil
-			}, 30*time.Second, time.Second).ShouldNot(HaveOccurred())
-		}
-
 		BeforeEach(func() {
 			nodeList = libnode.GetAllSchedulableNodes(virtCli).Items
 			for _, node := range nodeList {
-				setNodeUnschedulable(node.Name)
+				libnode.SetNodeUnschedulable(node.Name, virtCli)
 			}
 			eventuallyWithTimeout(waitForDeploymentsToStabilize)
 		})
 
 		AfterEach(func() {
 			for _, node := range nodeList {
-				setNodeSchedulable(node.Name)
+				libnode.SetNodeSchedulable(node.Name, virtCli)
 			}
 			eventuallyWithTimeout(waitForDeploymentsToStabilize)
 		})