daemonset pkg: WithPodAffinity method added

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 {

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{