From be02d00badfba81ff3e1d0dae5887d6d0faeba7f Mon Sep 17 00:00:00 2001
From: Itamar Holder <iholder@redhat.com>
Date: Thu, 11 Apr 2024 11:08:22 +0300
Subject: [PATCH] Refactor: generatePatchesForPodWithNodeName() to handle
 nodeName case separately

Signed-off-by: Itamar Holder <iholder@redhat.com>
---
 pkg/aaq-server/handler/handler.go | 86 +++++++++++++++++--------------
 1 file changed, 46 insertions(+), 40 deletions(-)

diff --git a/pkg/aaq-server/handler/handler.go b/pkg/aaq-server/handler/handler.go
index 7e739660..01995d41 100644
--- a/pkg/aaq-server/handler/handler.go
+++ b/pkg/aaq-server/handler/handler.go
@@ -85,46 +85,7 @@ func (v Handler) mutatePod() (*admissionv1.AdmissionReview, error) {
 		},
 	}
 
-	if pod.Spec.NodeName != "" {
-		affinity := pod.Spec.Affinity
-		affinityPatchOp := patch.PatchReplaceOp
-
-		if affinity == nil {
-			affinity = &v1.Affinity{}
-			affinityPatchOp = patch.PatchAddOp
-		}
-		if affinity.NodeAffinity == nil {
-			affinity.NodeAffinity = &v1.NodeAffinity{}
-		}
-		if affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil {
-			affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = &v1.NodeSelector{}
-		}
-		affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms = append(
-			affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms,
-			v1.NodeSelectorTerm{
-				MatchFields: []v1.NodeSelectorRequirement{
-					{
-						Key:      "metadata.name",
-						Operator: v1.NodeSelectorOpIn,
-						Values:   []string{pod.Spec.NodeName},
-					},
-				},
-			},
-		)
-
-		patches = append(patches,
-			patch.PatchOperation{
-				Op:    affinityPatchOp,
-				Path:  "/spec/affinity",
-				Value: affinity,
-			},
-			patch.PatchOperation{
-				Op:    patch.PatchReplaceOp,
-				Path:  "/spec/nodeName",
-				Value: "",
-			},
-		)
-	}
+	patches = append(patches, generatePatchesForPodWithNodeName(pod)...)
 
 	patchBytes, err := patch.GeneratePatchPayload(patches...)
 	if err != nil {
@@ -259,3 +220,48 @@ func getResourcesNames(resourceList v1.ResourceList) []v1.ResourceName {
 	}
 	return keys
 }
+
+func generatePatchesForPodWithNodeName(pod v1.Pod) []patch.PatchOperation {
+	if pod.Spec.NodeName == "" {
+		return nil
+	}
+
+	affinity := pod.Spec.Affinity
+	affinityPatchOp := patch.PatchReplaceOp
+
+	if affinity == nil {
+		affinity = &v1.Affinity{}
+		affinityPatchOp = patch.PatchAddOp
+	}
+	if affinity.NodeAffinity == nil {
+		affinity.NodeAffinity = &v1.NodeAffinity{}
+	}
+	if affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil {
+		affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = &v1.NodeSelector{}
+	}
+	affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms = append(
+		affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms,
+		v1.NodeSelectorTerm{
+			MatchFields: []v1.NodeSelectorRequirement{
+				{
+					Key:      "metadata.name",
+					Operator: v1.NodeSelectorOpIn,
+					Values:   []string{pod.Spec.NodeName},
+				},
+			},
+		},
+	)
+
+	return []patch.PatchOperation{
+		{
+			Op:    affinityPatchOp,
+			Path:  "/spec/affinity",
+			Value: affinity,
+		},
+		{
+			Op:    patch.PatchReplaceOp,
+			Path:  "/spec/nodeName",
+			Value: "",
+		},
+	}
+}