test: zone refactor

test/pkg/environment/aws/environment.go

@@ -82,6 +82,13 @@ type Environment struct {
 	ClusterEndpoint   string
 	InterruptionQueue string
 	PrivateCluster    bool
+	ZoneInfo          []ZoneInfo
+}
+
+type ZoneInfo struct {
+	Zone     string
+	ZoneID   string
+	ZoneType string
 }
 
 func NewEnvironment(t *testing.T) *Environment {
@@ -123,6 +130,14 @@ func NewEnvironment(t *testing.T) *Environment {
 		out := lo.Must(sqsapi.GetQueueUrlWithContext(env.Context, &servicesqs.GetQueueUrlInput{QueueName: aws.String(v)}))
 		awsEnv.SQSProvider = lo.Must(sqs.NewDefaultProvider(sqsapi, lo.FromPtr(out.QueueUrl)))
 	}
+	// Populate ZoneInfo for all AZs in the region
+	awsEnv.ZoneInfo = lo.Map(lo.Must(awsEnv.EC2API.DescribeAvailabilityZones(&ec2.DescribeAvailabilityZonesInput{})).AvailabilityZones, func(zone *ec2.AvailabilityZone, _ int) ZoneInfo {
+		return ZoneInfo{
+			Zone:     lo.FromPtr(zone.ZoneName),
+			ZoneID:   lo.FromPtr(zone.ZoneId),
+			ZoneType: lo.FromPtr(zone.ZoneType),
+		}
+	})
 	return awsEnv
 }
 

test/pkg/environment/aws/expectations.go

@@ -210,22 +210,6 @@ func (env *Environment) GetSpotInstanceRequest(id *string) *ec2.SpotInstanceRequ
 	return siro.SpotInstanceRequests[0]
 }
 
-// GetZones returns all available zones mapped from zone -> zone type
-func (env *Environment) GetZones() map[string]string {
-	output := lo.Must(env.EC2API.DescribeAvailabilityZones(&ec2.DescribeAvailabilityZonesInput{}))
-	return lo.Associate(output.AvailabilityZones, func(zone *ec2.AvailabilityZone) (string, string) {
-		return lo.FromPtr(zone.ZoneName), lo.FromPtr(zone.ZoneType)
-	})
-}
-
-// GetZoneIDMapping returns map from zone to zone id
-func (env *Environment) GetZoneIDMapping() map[string]string {
-	output := lo.Must(env.EC2API.DescribeAvailabilityZones(&ec2.DescribeAvailabilityZonesInput{}))
-	return lo.Associate(output.AvailabilityZones, func(zone *ec2.AvailabilityZone) (string, string) {
-		return lo.FromPtr(zone.ZoneName), lo.FromPtr(zone.ZoneId)
-	})
-}
-
 // GetSubnets returns all subnets matching the label selector
 // mapped from AZ -> {subnet-ids...}
 func (env *Environment) GetSubnets(tags map[string]string) map[string][]string {
@@ -251,10 +235,11 @@ func (env *Environment) GetSubnets(tags map[string]string) map[string][]string {
 type SubnetInfo struct {
 	Name string
 	ID   string
+	ZoneInfo
 }
 
-// GetSubnetNameAndIds returns all subnets matching the label selector
-func (env *Environment) GetSubnetNameAndIds(tags map[string]string) []SubnetInfo {
+// GetSubnetInfo returns all subnets matching the label selector
+func (env *Environment) GetSubnetInfo(tags map[string]string) []SubnetInfo {
 	var filters []*ec2.Filter
 	for key, val := range tags {
 		filters = append(filters, &ec2.Filter{
@@ -269,6 +254,11 @@ func (env *Environment) GetSubnetNameAndIds(tags map[string]string) []SubnetInfo
 			if tag, ok := lo.Find(s.Tags, func(t *ec2.Tag) bool { return aws.StringValue(t.Key) == "Name" }); ok {
 				elem.Name = aws.StringValue(tag.Value)
 			}
+			if info, ok := lo.Find(env.ZoneInfo, func(info ZoneInfo) bool {
+				return aws.StringValue(s.AvailabilityZone) == info.Zone
+			}); ok {
+				elem.ZoneInfo = info
+			}
 			return elem
 		})
 		return true

test/suites/drift/suite_test.go

@@ -515,7 +515,7 @@ var _ = Describe("Drift", func() {
 		env.EventuallyExpectHealthyPodCount(selector, numPods)
 	})
 	It("should disrupt nodes that have drifted due to subnets", func() {
-		subnets := env.GetSubnetNameAndIds(map[string]string{"karpenter.sh/discovery": env.ClusterName})
+		subnets := env.GetSubnetInfo(map[string]string{"karpenter.sh/discovery": env.ClusterName})
 		Expect(len(subnets)).To(BeNumerically(">", 1))
 
 		nodeClass.Spec.SubnetSelectorTerms = []v1beta1.SubnetSelectorTerm{{ID: subnets[0].ID}}

test/suites/integration/scheduling_test.go

@@ -115,12 +115,7 @@ var _ = Describe("Scheduling", Ordered, ContinueOnFailure, func() {
 					{
 						Key:      v1beta1.LabelTopologyZoneID,
 						Operator: v1.NodeSelectorOpIn,
-						Values: func() []string {
-							zoneIDMapping := env.GetZoneIDMapping()
-							subnetZones := lo.Keys(env.GetSubnets(map[string]string{"karpenter.sh/discovery": env.ClusterName}))
-							targetZoneID := zoneIDMapping[subnetZones[0]]
-							return []string{targetZoneID}
-						}(),
+						Values:   []string{env.GetSubnetInfo(map[string]string{"karpenter.sh/discovery": env.ClusterName})[0].ZoneInfo.ZoneID},
 					},
 				},
 			}})
@@ -603,17 +598,8 @@ var _ = Describe("Scheduling", Ordered, ContinueOnFailure, func() {
 		})
 
 		It("should provision a node for a pod with overlapping zone and zone-id requirements", func() {
-			type mapping struct {
-				zone   string
-				zoneID string
-			}
-			subnetZones := lo.Keys(env.GetSubnets(map[string]string{"karpenter.sh/discovery": env.ClusterName}))
-			zones := lo.Filter(lo.MapToSlice(env.GetZoneIDMapping(), func(zone, zoneID string) mapping {
-				return mapping{zone, zoneID}
-			}), func(m mapping, _ int) bool {
-				return lo.Contains(subnetZones, m.zone)
-			})
-			Expect(len(zones)).To(BeNumerically(">=", 3))
+			subnetInfo := env.GetSubnetInfo(map[string]string{"karpenter.sh/discovery": env.ClusterName})
+			Expect(len(subnetInfo)).To(BeNumerically(">=", 3))
 
 			// Create a pod with 'overlapping' zone and zone-id requirements. With two options for each label, but only one pair of zone-zoneID that maps to the
 			// same AZ, we will always expect the pod to be scheduled to that AZ. In this case, this is the mapping at zone[1].
@@ -622,19 +608,19 @@ var _ = Describe("Scheduling", Ordered, ContinueOnFailure, func() {
 					{
 						Key:      v1.LabelTopologyZone,
 						Operator: v1.NodeSelectorOpIn,
-						Values:   lo.Map(zones[0:2], func(m mapping, _ int) string { return m.zone }),
+						Values:   lo.Map(subnetInfo[0:2], func(info aws.SubnetInfo, _ int) string { return info.Zone }),
 					},
 					{
 						Key:      v1beta1.LabelTopologyZoneID,
 						Operator: v1.NodeSelectorOpIn,
-						Values:   lo.Map(zones[1:3], func(m mapping, _ int) string { return m.zoneID }),
+						Values:   lo.Map(subnetInfo[1:3], func(info aws.SubnetInfo, _ int) string { return info.Zone }),
 					},
 				},
 			})
 			env.ExpectCreated(nodePool, nodeClass, pod)
 			node := env.EventuallyExpectNodeCount("==", 1)[0]
-			Expect(node.Labels[v1.LabelTopologyZone]).To(Equal(zones[1].zone))
-			Expect(node.Labels[v1beta1.LabelTopologyZoneID]).To(Equal(zones[1].zoneID))
+			Expect(node.Labels[v1.LabelTopologyZone]).To(Equal(subnetInfo[1].Zone))
+			Expect(node.Labels[v1beta1.LabelTopologyZoneID]).To(Equal(subnetInfo[1].ZoneID))
 		})
 		It("should provision nodes for pods with zone-id requirements in the correct zone", func() {
 			const expectedZoneLabel = "domain-label"
@@ -645,20 +631,19 @@ var _ = Describe("Scheduling", Ordered, ContinueOnFailure, func() {
 				},
 			})
 
-			// zoneIDMapping is a mapping from zone to zone-id for all zoneIDMapping available in the provided subnets
-			zoneIDMapping := lo.PickByKeys(env.GetZoneIDMapping(), lo.Keys(env.GetSubnets(map[string]string{"karpenter.sh/discovery": env.ClusterName})))
-			pods := lo.MapToSlice(zoneIDMapping, func(zone, zoneID string) *v1.Pod {
+			subnetInfo := env.GetSubnetInfo(map[string]string{"karpenter.sh/discovery": env.ClusterName})
+			pods := lo.Map(subnetInfo, func(info aws.SubnetInfo, _ int) *v1.Pod {
 				return test.Pod(test.PodOptions{
 					NodeRequirements: []v1.NodeSelectorRequirement{
 						{
 							Key:      expectedZoneLabel,
 							Operator: v1.NodeSelectorOpIn,
-							Values:   []string{zone},
+							Values:   []string{info.Zone},
 						},
 						{
 							Key:      v1beta1.LabelTopologyZoneID,
 							Operator: v1.NodeSelectorOpIn,
-							Values:   []string{zoneID},
+							Values:   []string{info.ZoneID},
 						},
 					},
 				})
@@ -668,12 +653,16 @@ var _ = Describe("Scheduling", Ordered, ContinueOnFailure, func() {
 			for _, pod := range pods {
 				env.ExpectCreated(pod)
 			}
-			nodes := env.EventuallyExpectCreatedNodeCount("==", len(zoneIDMapping))
+			nodes := env.EventuallyExpectCreatedNodeCount("==", len(subnetInfo))
 			for _, node := range nodes {
 				expectedZone, ok := node.Labels[expectedZoneLabel]
 				Expect(ok).To(BeTrue())
 				Expect(node.Labels[v1.LabelTopologyZone]).To(Equal(expectedZone))
-				Expect(node.Labels[v1beta1.LabelTopologyZoneID]).To(Equal(zoneIDMapping[expectedZone]))
+				zoneInfo, ok := lo.Find(subnetInfo, func(info aws.SubnetInfo) bool {
+					return info.Zone == expectedZone
+				})
+				Expect(ok).To(BeTrue())
+				Expect(node.Labels[v1beta1.LabelTopologyZoneID]).To(Equal(zoneInfo.ZoneID))
 			}
 		})
 	})

test/suites/integration/subnet_test.go

@@ -78,7 +78,7 @@ var _ = Describe("Subnets", func() {
 	})
 	It("should use the subnet tag selector with multiple tag values", func() {
 		// Get all the subnets for the cluster
-		subnets := env.GetSubnetNameAndIds(map[string]string{"karpenter.sh/discovery": env.ClusterName})
+		subnets := env.GetSubnetInfo(map[string]string{"karpenter.sh/discovery": env.ClusterName})
 		Expect(len(subnets)).To(BeNumerically(">", 1))
 		firstSubnet := subnets[0]
 		lastSubnet := subnets[len(subnets)-1]

test/suites/localzone/suite_test.go

@@ -70,8 +70,11 @@ var _ = BeforeEach(func() {
 		NodeSelectorRequirement: v1.NodeSelectorRequirement{
 			Key:      v1.LabelTopologyZone,
 			Operator: v1.NodeSelectorOpIn,
-			Values:   lo.Keys(lo.PickByValues(env.GetZones(), []string{"local-zone"})),
-		}})
+			Values: lo.FilterMap(env.GetSubnetInfo(map[string]string{"karpenter.sh/discovery": env.ClusterName}), func(info aws.SubnetInfo, _ int) (string, bool) {
+				return info.Zone, info.ZoneType == "local-zone"
+			}),
+		},
+	})
 })
 var _ = AfterEach(func() { env.Cleanup() })
 var _ = AfterEach(func() { env.AfterEach() })

test/suites/nodeclaim/garbage_collection_test.go

@@ -45,7 +45,7 @@ var _ = Describe("GarbageCollection", func() {
 
 	BeforeEach(func() {
 		securityGroups := env.GetSecurityGroups(map[string]string{"karpenter.sh/discovery": env.ClusterName})
-		subnets := env.GetSubnetNameAndIds(map[string]string{"karpenter.sh/discovery": env.ClusterName})
+		subnets := env.GetSubnetInfo(map[string]string{"karpenter.sh/discovery": env.ClusterName})
 		Expect(securityGroups).ToNot(HaveLen(0))
 		Expect(subnets).ToNot(HaveLen(0))