Fix AZ Tests by limit to Cluster AZs

test/framework/resources/aws/services/ec2.go

@@ -36,7 +36,7 @@ type EC2 interface {
 	AssociateVPCCIDRBlock(ctx context.Context, vpcId string, cidrBlock string) (*ec2.AssociateVpcCidrBlockOutput, error)
 	TerminateInstance(ctx context.Context, instanceIDs []string) error
 	DisAssociateVPCCIDRBlock(ctx context.Context, associationID string) error
-	DescribeSubnet(ctx context.Context, subnetID string) (*ec2.DescribeSubnetsOutput, error)
+	DescribeSubnets(ctx context.Context, subnetIDs []string) (*ec2.DescribeSubnetsOutput, error)
 	CreateSubnet(ctx context.Context, cidrBlock string, vpcID string, az string) (*ec2.CreateSubnetOutput, error)
 	DeleteSubnet(ctx context.Context, subnetID string) error
 	DescribeRouteTables(ctx context.Context, subnetID string) (*ec2.DescribeRouteTablesOutput, error)
@@ -290,9 +290,9 @@ func (d *defaultEC2) CreateSubnet(ctx context.Context, cidrBlock string, vpcID s
 	return d.client.CreateSubnet(ctx, createSubnetInput)
 }
 
-func (d *defaultEC2) DescribeSubnet(ctx context.Context, subnetID string) (*ec2.DescribeSubnetsOutput, error) {
+func (d *defaultEC2) DescribeSubnets(ctx context.Context, subnetIDs []string) (*ec2.DescribeSubnetsOutput, error) {
 	describeSubnetInput := &ec2.DescribeSubnetsInput{
-		SubnetIds: []string{subnetID},
+		SubnetIds: subnetIDs,
 	}
 	return d.client.DescribeSubnets(ctx, describeSubnetInput)
 }

test/framework/resources/aws/utils/nodegroup.go

@@ -281,7 +281,7 @@ func GetClusterVPCConfig(f *framework.Framework) (*ClusterVPCConfig, error) {
 
 	uniqueAZ := map[string]bool{}
 	for _, subnet := range clusterConfig.PublicSubnetList {
-		describeSubnet, err := f.CloudServices.EC2().DescribeSubnet(context.TODO(), subnet)
+		describeSubnet, err := f.CloudServices.EC2().DescribeSubnets(context.TODO(), []string{subnet})
 		if err != nil {
 			return nil, fmt.Errorf("failed to describe the subnet %s: %v", subnet, err)
 		}

test/integration/az-traffic/pod_traffic_across_az_test.go

@@ -28,6 +28,45 @@ const MetricNamespace = "NetworkingAZConnectivity"
 
 var f *framework.Framework
 
+func getClusterAZs(ctx context.Context, clusterName string) ([]string, error) {
+	// Step 1: Get all AZs in the region
+	describeAZOutput, err := f.CloudServices.EC2().DescribeAvailabilityZones(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to describe availability zones: %w", err)
+	}
+
+	// Step 2: Get cluster info
+	clusterInfo, err := f.CloudServices.EKS().DescribeCluster(ctx, clusterName)
+	if err != nil {
+		return nil, fmt.Errorf("failed to describe cluster: %w", err)
+	}
+
+	// Step 3: Get subnet info
+	subnetIDs := clusterInfo.Cluster.ResourcesVpcConfig.SubnetIds
+	describeSubnetsOutput, err := f.CloudServices.EC2().DescribeSubnets(ctx, subnetIDs)
+	if err != nil {
+		return nil, fmt.Errorf("failed to describe subnets: %w", err)
+	}
+
+	// Step 4: Create a map of cluster AZs
+	clusterAZs := make(map[string]bool)
+	for _, subnet := range describeSubnetsOutput.Subnets {
+		clusterAZs[*subnet.AvailabilityZone] = true
+	}
+
+	// Step 5: Filter and deduplicate AZs
+	var filteredAZs []string
+	azSet := make(map[string]bool)
+	for _, az := range describeAZOutput.AvailabilityZones {
+		if clusterAZs[*az.ZoneName] && !azSet[*az.ZoneName] {
+			filteredAZs = append(filteredAZs, *az.ZoneName)
+			azSet[*az.ZoneName] = true
+		}
+	}
+
+	return filteredAZs, nil
+}
+
 var _ = Describe("[STATIC_CANARY] AZ Node Presence", FlakeAttempts(retries), func() {
 
 	Context("While testing AZ availability", func() {
@@ -37,11 +76,11 @@ var _ = Describe("[STATIC_CANARY] AZ Node Presence", FlakeAttempts(retries), fun
 			nodes, err := f.K8sResourceManagers.NodeManager().GetNodes(f.Options.NgNameLabelKey, f.Options.NgNameLabelVal)
 			Expect(err).ToNot(HaveOccurred())
 
-			describeAZOutput, err := f.CloudServices.EC2().DescribeAvailabilityZones(context.TODO())
+			clusterAZs, err := getClusterAZs(context.TODO(), f.Options.ClusterName)
 			Expect(err).ToNot(HaveOccurred())
 
-			for _, az := range describeAZOutput.AvailabilityZones {
-				nodePresenceInAZ[*az.ZoneName] = false
+			for _, az := range clusterAZs {
+				nodePresenceInAZ[az] = false
 			}
 			for i := range nodes.Items {
 				// node label key "topology.kubernetes.io/zone" is well known label populated by cloud controller manager

test/integration/eni-subnet-discovery/eni_subnet_discovery_test.go

@@ -319,7 +319,7 @@ func checkSecondaryENISubnets(expectNewCidr bool) {
 
 	By(fmt.Sprintf("checking the secondary ENI subnets are in the CIDR %s", expectedCidrRangeString))
 	for _, subnetID := range newEniSubnetIds {
-		subnetOutput, err := f.CloudServices.EC2().DescribeSubnet(context.TODO(), subnetID)
+		subnetOutput, err := f.CloudServices.EC2().DescribeSubnets(context.TODO(), []string{subnetID})
 		Expect(err).ToNot(HaveOccurred())
 		cidrSplit := strings.Split(*subnetOutput.Subnets[0].CidrBlock, "/")
 		actualSubnetIp, _, _ := net.ParseCIDR(*subnetOutput.Subnets[0].CidrBlock)

test/integration/snat/snat_test.go

@@ -73,7 +73,7 @@ var _ = Describe("SNAT tests", func() {
 	Context("Validate AWS_VPC_K8S_CNI_EXCLUDE_SNAT_CIDRS", func() {
 		It("Verify External Domain Connectivity by modifying AWS_VPC_K8S_CNI_EXCLUDE_SNAT_CIDRS", func() {
 			By("Getting CIDR for primary node's private subnet")
-			out, err := f.CloudServices.EC2().DescribeSubnet(context.TODO(), privateSubnetId)
+			out, err := f.CloudServices.EC2().DescribeSubnets(context.TODO(), []string{privateSubnetId})
 			Expect(err).NotTo(HaveOccurred())
 			Expect(len(out.Subnets)).To(BeNumerically(">", 0))