From b23a605061fabfa6eb08a4cc0223c04c0af090fc Mon Sep 17 00:00:00 2001
From: Jason Deal <jmdeal@amazon.com>
Date: Mon, 11 Dec 2023 13:27:53 -0800
Subject: [PATCH] update tests

---
 .../provisioning/scheduling/topology_test.go  | 115 +++++++++++++++++-
 1 file changed, 114 insertions(+), 1 deletion(-)

diff --git a/pkg/controllers/provisioning/scheduling/topology_test.go b/pkg/controllers/provisioning/scheduling/topology_test.go
index 9db43a3a52..b92ae41d1f 100644
--- a/pkg/controllers/provisioning/scheduling/topology_test.go
+++ b/pkg/controllers/provisioning/scheduling/topology_test.go
@@ -1319,9 +1319,123 @@ var _ = Describe("Topology", func() {
 			// and should schedule all of the pods on the same node
 			ExpectSkew(ctx, env.Client, "default", &topology[0]).To(ConsistOf(5))
 		})
+		It("should balance pods across a label when discovered from the provisioner (NodeTaintsPolicy=ignore)", func() {
+			const spreadLabel = "fake-label"
+			const taintKey = "taint-key"
+			nodePool.Spec.Template.Spec.Requirements = append(nodePool.Spec.Template.Spec.Requirements, v1.NodeSelectorRequirement{
+				Key:      spreadLabel,
+				Operator: v1.NodeSelectorOpIn,
+				Values:   []string{"foo"},
+			})
+			taintedNodePool := test.NodePool(v1beta1.NodePool{
+				Spec: v1beta1.NodePoolSpec{
+					Template: v1beta1.NodeClaimTemplate{
+						Spec: v1beta1.NodeClaimSpec{
+							Taints: []v1.Taint{
+								{
+									Key:    taintKey,
+									Value:  "taint-value",
+									Effect: v1.TaintEffectNoSchedule,
+								},
+							},
+							Requirements: []v1.NodeSelectorRequirement{
+								{
+									Key:      v1beta1.CapacityTypeLabelKey,
+									Operator: v1.NodeSelectorOpExists,
+								},
+								{
+									Key:      spreadLabel,
+									Operator: v1.NodeSelectorOpIn,
+									Values:   []string{"bar"},
+								},
+							},
+						},
+					},
+				},
+			})
+
+			honor := v1.NodeInclusionPolicyIgnore
+			topology := []v1.TopologySpreadConstraint{{
+				TopologyKey:       spreadLabel,
+				WhenUnsatisfiable: v1.DoNotSchedule,
+				LabelSelector:     &metav1.LabelSelector{MatchLabels: labels},
+				MaxSkew:           1,
+				NodeTaintsPolicy:  &honor,
+			}}
+
+			pods := test.UnschedulablePods(test.PodOptions{
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: labels,
+				},
+				TopologySpreadConstraints: topology,
+			}, 2)
+
+			ExpectApplied(ctx, env.Client, nodePool, taintedNodePool)
+			ExpectProvisioned(ctx, env.Client, cluster, cloudProvider, prov, pods...)
+
+			// should fail to schedule both pods, one pod is scheduled to domain "foo" but the other can't be scheduled to domain "bar"
+			ExpectSkew(ctx, env.Client, "default", &topology[0]).To(ConsistOf(1))
+		})
 		It("should balance pods across a label when discovered from the provisioner (NodeTaintsPolicy=honor)", func() {
 			const spreadLabel = "fake-label"
 			const taintKey = "taint-key"
+			nodePool.Spec.Template.Spec.Requirements = append(nodePool.Spec.Template.Spec.Requirements, v1.NodeSelectorRequirement{
+				Key:      spreadLabel,
+				Operator: v1.NodeSelectorOpIn,
+				Values:   []string{"foo"},
+			})
+			taintedNodePool := test.NodePool(v1beta1.NodePool{
+				Spec: v1beta1.NodePoolSpec{
+					Template: v1beta1.NodeClaimTemplate{
+						Spec: v1beta1.NodeClaimSpec{
+							Taints: []v1.Taint{
+								{
+									Key:    taintKey,
+									Value:  "taint-value",
+									Effect: v1.TaintEffectNoSchedule,
+								},
+							},
+							Requirements: []v1.NodeSelectorRequirement{
+								{
+									Key:      v1beta1.CapacityTypeLabelKey,
+									Operator: v1.NodeSelectorOpExists,
+								},
+								{
+									Key:      spreadLabel,
+									Operator: v1.NodeSelectorOpIn,
+									Values:   []string{"bar"},
+								},
+							},
+						},
+					},
+				},
+			})
+
+			honor := v1.NodeInclusionPolicyHonor
+			topology := []v1.TopologySpreadConstraint{{
+				TopologyKey:       spreadLabel,
+				WhenUnsatisfiable: v1.DoNotSchedule,
+				LabelSelector:     &metav1.LabelSelector{MatchLabels: labels},
+				MaxSkew:           1,
+				NodeTaintsPolicy:  &honor,
+			}}
+
+			pods := test.UnschedulablePods(test.PodOptions{
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: labels,
+				},
+				TopologySpreadConstraints: topology,
+			}, 2)
+
+			ExpectApplied(ctx, env.Client, nodePool, taintedNodePool)
+			ExpectProvisioned(ctx, env.Client, cluster, cloudProvider, prov, pods...)
+
+			// should schedule all pods to domain "foo", ignoring bar since pods don't tolerate
+			ExpectSkew(ctx, env.Client, "default", &topology[0]).To(ConsistOf(2))
+		})
+		It("should balance pods across a label when mutually exclusive NodePools (by taints) share domains (NodeTaintsPolicy=honor)", func() {
+			const spreadLabel = "fake-label"
+			const taintKey = "taint-key"
 
 			nodePools := lo.Map([][]string{{"foo", "bar"}, {"foo", "baz"}}, func(domains []string, i int) *v1beta1.NodePool {
 				return test.NodePool(v1beta1.NodePool{
@@ -1377,7 +1491,6 @@ var _ = Describe("Topology", func() {
 
 			ExpectApplied(ctx, env.Client, nodePools[0], nodePools[1])
 			ExpectProvisioned(ctx, env.Client, cluster, cloudProvider, prov, pods...)
-			ExpectScheduled(ctx, env.Client, pods[0])
 
 			// Expect 3 total nodes provisioned, 2 pods schedule to foo, 1 to bar, and 1 to baz
 			ExpectSkew(ctx, env.Client, "default", &topology[0]).To(ConsistOf(1, 2, 1))