From 8138144d0ac7f94db2ac948de8dcb72f9efdf429 Mon Sep 17 00:00:00 2001
From: Elena German <elgerman@elgerman-thinkpadp16vgen1.rmtcaon.csb>
Date: Thu, 14 Nov 2024 13:26:52 -0500
Subject: [PATCH] daemonset pkg: WithPodAffinity method added

---
 pkg/daemonset/daemonset.go      | 14 ++++++++++++++
 pkg/daemonset/daemonset_test.go | 24 ++++++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/pkg/daemonset/daemonset.go b/pkg/daemonset/daemonset.go
index a8d0c07bd..eb319283c 100644
--- a/pkg/daemonset/daemonset.go
+++ b/pkg/daemonset/daemonset.go
@@ -174,6 +174,20 @@ func (builder *Builder) WithHostNetwork() *Builder {
 	return builder
 }
 
+// WithPodAffinity applies pod's Affinity to daemonset definition.
+func (builder *Builder) WithPodAffinity(podAffinity *corev1.Affinity) *Builder {
+	if valid, _ := builder.validate(); !valid {
+		return builder
+	}
+
+	glog.V(100).Infof("Definding pod affinity to daemonset %s in namespace %s",
+		builder.Definition.Name, builder.Definition.Namespace)
+
+	builder.Definition.Spec.Template.Spec.Affinity = podAffinity
+
+	return builder
+}
+
 // WithVolume defines Volume of daemonset under PodTemplateSpec.
 func (builder *Builder) WithVolume(dsVolume corev1.Volume) *Builder {
 	if valid, _ := builder.validate(); !valid {
diff --git a/pkg/daemonset/daemonset_test.go b/pkg/daemonset/daemonset_test.go
index ade358da0..e9c104cf8 100644
--- a/pkg/daemonset/daemonset_test.go
+++ b/pkg/daemonset/daemonset_test.go
@@ -233,6 +233,30 @@ func TestDaemonsetWithHostNetwork(t *testing.T) {
 	assert.Equal(t, true, testBuilder.Definition.Spec.Template.Spec.HostNetwork)
 }
 
+func TestDaemonsetWithPodAffinity(t *testing.T) {
+	testBuilder := buildValidTestBuilderWithClient([]runtime.Object{})
+
+	podAffinity := corev1.Affinity{
+		NodeAffinity: &corev1.NodeAffinity{
+			RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
+				NodeSelectorTerms: []corev1.NodeSelectorTerm{
+					{
+						MatchExpressions: []corev1.NodeSelectorRequirement{
+							{
+								Key:      "kubernetes.io/hostname",
+								Operator: corev1.NodeSelectorOpIn,
+								Values:   []string{"host1", "host2"},
+							},
+						},
+					},
+				}},
+		},
+	}
+	testBuilder.WithPodAffinity(&podAffinity)
+
+	assert.Equal(t, &podAffinity, testBuilder.Definition.Spec.Template.Spec.Affinity)
+}
+
 func TestDaemonsetWithVolume(t *testing.T) {
 	testBuilder := buildValidTestBuilderWithClient([]runtime.Object{})
 	testBuilder.WithVolume(corev1.Volume{