repo_description: The official CLI for Amazon EKS
repo_topics: amazon-web-services aws-cloudformation aws-ec2 aws-eks eks kubernetes kubernetes-cluster kubernetes-deployment kubernetes-setup
code: package az_test
  
  import (
  	"context"
  	"fmt"
  
  	"github.com/aws/aws-sdk-go-v2/aws"
  	"github.com/aws/aws-sdk-go-v2/service/ec2"
  	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
  	. "github.com/onsi/ginkgo/v2"
  	. "github.com/onsi/gomega"
  	"github.com/stretchr/testify/mock"
  
  	api "github.com/weaveworks/eksctl/pkg/apis/eksctl.io/v1alpha5"
  	"github.com/weaveworks/eksctl/pkg/az"
  	"github.com/weaveworks/eksctl/pkg/testutils/mockprovider"
  )
  
  var _ = Describe("AZ", func() {
  	var (
  		region string
  		p      *mockprovider.MockProvider
  		spec   *api.ClusterConfig
  	)
  
  	BeforeEach(func() {
  		region = "us-west-1"
  		p = mockprovider.NewMockProvider()
  		spec = api.NewClusterConfig()
  	})
  
  	When("1 AZ is available", func() {
  		BeforeEach(func() {
  			p.MockEC2().On("DescribeAvailabilityZones", mock.Anything, &ec2.DescribeAvailabilityZonesInput{
  				Filters: []ec2types.Filter{
  					{
  						Name:   aws.String("region-name"),
  						Values: []string{region},
  					},
  					{
  						Name:   aws.String("state"),
  						Values: []string{string(ec2types.AvailabilityZoneStateAvailable)},
  					},
  					{
  						Name:   aws.String("zone-type"),
  						Values: []string{string(ec2types.LocationTypeAvailabilityZone)},
  					},
  				},
  			}).Return(&ec2.DescribeAvailabilityZonesOutput{
  				AvailabilityZones: []ec2types.AvailabilityZone{
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone1"),
  				},
  			}, nil)
  		})
  
  		It("errors", func() {
  			_, err := az.GetAvailabilityZones(context.Background(), p.MockEC2(), region, spec)
  			Expect(err).To(MatchError("only 1 zones discovered [zone1], at least 2 are required"))
  		})
  	})
  
  	When("2 AZs are available", func() {
  		BeforeEach(func() {
  			p.MockEC2().On("DescribeAvailabilityZones", mock.Anything, &ec2.DescribeAvailabilityZonesInput{
  				Filters: []ec2types.Filter{
  					{
  						Name:   aws.String("region-name"),
  						Values: []string{region},
  					},
  					{
  						Name:   aws.String("state"),
  						Values: []string{string(ec2types.AvailabilityZoneStateAvailable)},
  					},
  					{
  						Name:   aws.String("zone-type"),
  						Values: []string{string(ec2types.LocationTypeAvailabilityZone)},
  					},
  				},
  			}).Return(&ec2.DescribeAvailabilityZonesOutput{
  				AvailabilityZones: []ec2types.AvailabilityZone{
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone1"),
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone2"),
  				},
  			}, nil)
  		})
  
  		It("should return the 2 available AZs", func() {
  			zones, err := az.GetAvailabilityZones(context.Background(), p.MockEC2(), region, spec)
  			Expect(err).NotTo(HaveOccurred())
  			Expect(zones).To(HaveLen(2))
  			Expect(zones).To(ConsistOf("zone1", "zone2"))
  		})
  	})
  
  	When("3 AZs are available", func() {
  		BeforeEach(func() {
  			p.MockEC2().On("DescribeAvailabilityZones", mock.Anything, &ec2.DescribeAvailabilityZonesInput{
  				Filters: []ec2types.Filter{
  					{
  						Name:   aws.String("region-name"),
  						Values: []string{region},
  					},
  					{
  						Name:   aws.String("state"),
  						Values: []string{string(ec2types.AvailabilityZoneStateAvailable)},
  					},
  					{
  						Name:   aws.String("zone-type"),
  						Values: []string{string(ec2types.LocationTypeAvailabilityZone)},
  					},
  				},
  			}).Return(&ec2.DescribeAvailabilityZonesOutput{
  				AvailabilityZones: []ec2types.AvailabilityZone{
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone1"),
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone2"),
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone3"),
  				},
  			}, nil)
  		})
  
  		It("should return the 3 available AZs", func() {
  			zones, err := az.GetAvailabilityZones(context.Background(), p.MockEC2(), region, spec)
  			Expect(err).NotTo(HaveOccurred())
  			Expect(zones).To(HaveLen(3))
  			Expect(zones).To(ConsistOf("zone1", "zone2", "zone3"))
  		})
  	})
  
  	When("more than 3 AZs are available", func() {
  		BeforeEach(func() {
  			p.MockEC2().On("DescribeAvailabilityZones", mock.Anything, &ec2.DescribeAvailabilityZonesInput{
  				Filters: []ec2types.Filter{
  					{
  						Name:   aws.String("region-name"),
  						Values: []string{region},
  					},
  					{
  						Name:   aws.String("state"),
  						Values: []string{string(ec2types.AvailabilityZoneStateAvailable)},
  					},
  					{
  						Name:   aws.String("zone-type"),
  						Values: []string{string(ec2types.LocationTypeAvailabilityZone)},
  					},
  				},
  			}).Return(&ec2.DescribeAvailabilityZonesOutput{
  				AvailabilityZones: []ec2types.AvailabilityZone{
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone1"),
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone2"),
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone3"),
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone4"),
  				},
  			}, nil)
  		})
  
  		It("should return a random set of 3 available AZs", func() {
  			zones, err := az.GetAvailabilityZones(context.Background(), p.MockEC2(), region, spec)
  			Expect(err).NotTo(HaveOccurred())
  			Expect(zones).To(HaveLen(3))
  			Expect(zonesAreUnique(zones)).To(BeTrue())
  		})
  	})
  
  	When("instance types are defined", func() {
  		BeforeEach(func() {
  			spec = api.NewClusterConfig()
  			spec.NodeGroups = []*api.NodeGroup{
  				{
  					NodeGroupBase: &api.NodeGroupBase{
  						Name:         "test-az-1",
  						InstanceType: "t2.small",
  					},
  				},
  				{
  					NodeGroupBase: &api.NodeGroupBase{
  						Name:         "test-az-2",
  						InstanceType: "t2.medium",
  					},
  				},
  			}
  			p.MockEC2().On("DescribeAvailabilityZones", mock.Anything, &ec2.DescribeAvailabilityZonesInput{
  				Filters: []ec2types.Filter{
  					{
  						Name:   aws.String("region-name"),
  						Values: []string{region},
  					},
  					{
  						Name:   aws.String("state"),
  						Values: []string{string(ec2types.AvailabilityZoneStateAvailable)},
  					},
  					{
  						Name:   aws.String("zone-type"),
  						Values: []string{string(ec2types.LocationTypeAvailabilityZone)},
  					},
  				},
  			}).Return(&ec2.DescribeAvailabilityZonesOutput{
  				AvailabilityZones: []ec2types.AvailabilityZone{
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone1"),
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone2"),
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone3"),
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone4"),
  				},
  			}, nil)
  			// split DescribeInstanceTypeOfferings response in two pages so we unit test the use of the paginator at the same time
  			p.MockEC2().On("DescribeInstanceTypeOfferings", mock.Anything, &ec2.DescribeInstanceTypeOfferingsInput{
  				Filters: []ec2types.Filter{
  					{
  						Name:   aws.String("instance-type"),
  						Values: []string{"t2.small", "t2.medium"},
  					},
  					{
  						Name:   aws.String("location"),
  						Values: []string{"zone1", "zone2", "zone3", "zone4"},
  					},
  				},
  				LocationType: ec2types.LocationTypeAvailabilityZone,
  				MaxResults:   aws.Int32(100),
  			}, mock.Anything).Return(&ec2.DescribeInstanceTypeOfferingsOutput{
  				NextToken: aws.String("token"),
  				InstanceTypeOfferings: []ec2types.InstanceTypeOffering{
  					{
  						InstanceType: "t2.small",
  						Location:     aws.String("zone1"),
  						LocationType: "availability-zone",
  					},
  					{
  						InstanceType: "t2.small",
  						Location:     aws.String("zone2"),
  						LocationType: "availability-zone",
  					},
  					{
  						InstanceType: "t2.small",
  						Location:     aws.String("zone4"),
  						LocationType: "availability-zone",
  					},
  				},
  			}, nil)
  			p.MockEC2().On("DescribeInstanceTypeOfferings", mock.Anything, &ec2.DescribeInstanceTypeOfferingsInput{
  				Filters: []ec2types.Filter{
  					{
  						Name:   aws.String("instance-type"),
  						Values: []string{"t2.small", "t2.medium"},
  					},
  					{
  						Name:   aws.String("location"),
  						Values: []string{"zone1", "zone2", "zone3", "zone4"},
  					},
  				},
  				LocationType: ec2types.LocationTypeAvailabilityZone,
  				MaxResults:   aws.Int32(100),
  				NextToken:    aws.String("token"),
  			}, mock.Anything).Return(&ec2.DescribeInstanceTypeOfferingsOutput{
  				InstanceTypeOfferings: []ec2types.InstanceTypeOffering{
  					{
  						InstanceType: "t2.medium",
  						Location:     aws.String("zone1"),
  						LocationType: "availability-zone",
  					},
  					{
  						InstanceType: "t2.medium",
  						Location:     aws.String("zone2"),
  						LocationType: "availability-zone",
  					},
  					{
  						InstanceType: "t2.medium",
  						Location:     aws.String("zone4"),
  						LocationType: "availability-zone",
  					},
  				},
  			}, nil)
  		})
  
  		It("should return only the zones which support all selected instance types", func() {
  			zones, err := az.GetAvailabilityZones(context.Background(), p.MockEC2(), region, spec)
  			Expect(err).NotTo(HaveOccurred())
  			Expect(zones).To(ConsistOf([]string{"zone1", "zone2", "zone4"}))
  			Expect(zonesAreUnique(zones)).To(BeTrue())
  		})
  	})
  
  	When("fetching the AZs errors", func() {
  		BeforeEach(func() {
  			p.MockEC2().On("DescribeAvailabilityZones", mock.Anything, &ec2.DescribeAvailabilityZonesInput{
  				Filters: []ec2types.Filter{
  					{
  						Name:   aws.String("region-name"),
  						Values: []string{region},
  					},
  					{
  						Name:   aws.String("state"),
  						Values: []string{string(ec2types.AvailabilityZoneStateAvailable)},
  					},
  					{
  						Name:   aws.String("zone-type"),
  						Values: []string{string(ec2types.LocationTypeAvailabilityZone)},
  					},
  				},
  			}).Return(&ec2.DescribeAvailabilityZonesOutput{}, fmt.Errorf("foo"))
  		})
  
  		It("errors", func() {
  			_, err := az.GetAvailabilityZones(context.Background(), p.MockEC2(), region, spec)
  			Expect(err).To(MatchError(fmt.Sprintf("error getting availability zones for region %s: foo", region)))
  		})
  	})
  
  	type unsupportedZoneEntry struct {
  		region        string
  		zoneNameToIDs map[string]string
  		expectedZones []string
  	}
  	DescribeTable("region with unsupported zone IDs", func(e unsupportedZoneEntry) {
  		var azs []ec2types.AvailabilityZone
  		for zoneName, zoneID := range e.zoneNameToIDs {
  			azs = append(azs, createAvailabilityZoneWithID(e.region, ec2types.AvailabilityZoneStateAvailable, zoneName, zoneID))
  		}
  		mockProvider := mockprovider.NewMockProvider()
  		mockProvider.MockEC2().On("DescribeAvailabilityZones", mock.Anything, &ec2.DescribeAvailabilityZonesInput{
  			Filters: []ec2types.Filter{
  				{
  					Name:   aws.String("region-name"),
  					Values: []string{e.region},
  				},
  				{
  					Name:   aws.String("state"),
  					Values: []string{string(ec2types.AvailabilityZoneStateAvailable)},
  				},
  				{
  					Name:   aws.String("zone-type"),
  					Values: []string{string(ec2types.LocationTypeAvailabilityZone)},
  				},
  			},
  		}).Return(&ec2.DescribeAvailabilityZonesOutput{
  			AvailabilityZones: azs,
  		}, nil)
  		mockProvider.MockEC2().On("DescribeInstanceTypeOfferings", mock.Anything, &ec2.DescribeInstanceTypeOfferingsInput{
  			Filters: []ec2types.Filter{
  				{
  					Name:   aws.String("instance-type"),
  					Values: []string{"t2.small", "t2.medium"},
  				},
  				{
  					Name:   aws.String("location"),
  					Values: []string{"zone1", "zone2", "zone3", "zone4"},
  				},
  			},
  			LocationType: ec2types.LocationTypeAvailabilityZone,
  			MaxResults:   aws.Int32(100),
  		}, mock.Anything).Return(&ec2.DescribeInstanceTypeOfferingsOutput{
  			NextToken: aws.String("token"),
  			InstanceTypeOfferings: []ec2types.InstanceTypeOffering{
  				{
  					InstanceType: "t2.small",
  					Location:     aws.String("zone1"),
  					LocationType: "availability-zone",
  				},
  				{
  					InstanceType: "t2.small",
  					Location:     aws.String("zone2"),
  					LocationType: "availability-zone",
  				},
  				{
  					InstanceType: "t2.small",
  					Location:     aws.String("zone4"),
  					LocationType: "availability-zone",
  				},
  				{
  					InstanceType: "t2.small",
  					Location:     aws.String("zone3"),
  					LocationType: "availability-zone",
  				},
  			},
  		}, nil)
  		clusterConfig := api.NewClusterConfig()
  		clusterConfig.Metadata.Region = e.region
  		clusterConfig.NodeGroups = []*api.NodeGroup{
  			{
  				NodeGroupBase: &api.NodeGroupBase{
  					Name: "test-az-1",
  				},
  			},
  			{
  				NodeGroupBase: &api.NodeGroupBase{
  					Name: "test-az-2",
  				},
  			},
  		}
  		zones, err := az.GetAvailabilityZones(context.Background(), mockProvider.MockEC2(), e.region, clusterConfig)
  		Expect(err).NotTo(HaveOccurred())
  		Expect(zones).To(ConsistOf(e.expectedZones))
  	},
  		Entry(api.RegionUSEast1, unsupportedZoneEntry{
  			region: api.RegionUSEast1,
  			zoneNameToIDs: map[string]string{
  				"zone1": "use1-az1",
  				"zone2": "use1-az3",
  				"zone3": "use1-az2",
  			},
  			expectedZones: []string{"zone1", "zone3"},
  		}),
  		Entry(api.RegionUSWest1, unsupportedZoneEntry{
  			region: api.RegionUSWest1,
  			zoneNameToIDs: map[string]string{
  				"zone1": "usw1-az2",
  				"zone2": "usw1-az1",
  				"zone3": "usw1-az3",
  			},
  			expectedZones: []string{"zone2", "zone3"},
  		}),
  		Entry(api.RegionCACentral1, unsupportedZoneEntry{
  			region: api.RegionCACentral1,
  			zoneNameToIDs: map[string]string{
  				"zone1": "cac1-az1",
  				"zone2": "cac1-az2",
  				"zone3": "cac1-az3",
  			},
  			expectedZones: []string{"zone1", "zone2"},
  		}),
  	)
  	When("the region contains zones that are denylisted", func() {
  		BeforeEach(func() {
  			region = api.RegionUSEast1
  
  			p.MockEC2().On("DescribeAvailabilityZones", mock.Anything, &ec2.DescribeAvailabilityZonesInput{
  				Filters: []ec2types.Filter{
  					{
  						Name:   aws.String("region-name"),
  						Values: []string{region},
  					},
  					{
  						Name:   aws.String("state"),
  						Values: []string{string(ec2types.AvailabilityZoneStateAvailable)},
  					},
  					{
  						Name:   aws.String("zone-type"),
  						Values: []string{string(ec2types.LocationTypeAvailabilityZone)},
  					},
  				},
  			}).Return(&ec2.DescribeAvailabilityZonesOutput{
  				AvailabilityZones: []ec2types.AvailabilityZone{
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone1"),
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone2"),
  					createAvailabilityZoneWithID(region, ec2types.AvailabilityZoneStateAvailable, "zone3", "use1-az3"),
  				},
  			}, nil)
  		})
  
  		It("should not use the denylisted zones", func() {
  			zones, err := az.GetAvailabilityZones(context.Background(), p.MockEC2(), region, spec)
  			Expect(err).NotTo(HaveOccurred())
  			Expect(zones).To(HaveLen(2))
  			Expect(zones).To(ConsistOf("zone1", "zone2"))
  		})
  	})
  
  	When("using us-east-1", func() {
  		BeforeEach(func() {
  			region = "us-east-1"
  
  			p.MockEC2().On("DescribeAvailabilityZones", mock.Anything, &ec2.DescribeAvailabilityZonesInput{
  				Filters: []ec2types.Filter{
  					{
  						Name:   aws.String("region-name"),
  						Values: []string{region},
  					},
  					{
  						Name:   aws.String("state"),
  						Values: []string{string(ec2types.AvailabilityZoneStateAvailable)},
  					},
  					{
  						Name:   aws.String("zone-type"),
  						Values: []string{string(ec2types.LocationTypeAvailabilityZone)},
  					},
  				},
  			}).Return(&ec2.DescribeAvailabilityZonesOutput{
  				AvailabilityZones: []ec2types.AvailabilityZone{
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone1"),
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone2"),
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone3"),
  					createAvailabilityZone(region, ec2types.AvailabilityZoneStateAvailable, "zone4"),
  				},
  			}, nil)
  		})
  
  		It("should only use 2 AZs, rather than the default 3", func() {
  			zones, err := az.GetAvailabilityZones(context.Background(), p.MockEC2(), region, spec)
  			Expect(err).NotTo(HaveOccurred())
  			Expect(zones).To(HaveLen(2))
  			Expect(zonesAreUnique(zones)).To(BeTrue())
  		})
  	})
  })
  
  func zonesAreUnique(zones []string) bool {
  	mapZones := make(map[string]interface{})
  	for _, z := range zones {
  		mapZones[z] = nil
  	}
  	return len(mapZones) == len(zones)
  }
  
  func createAvailabilityZone(region string, state ec2types.AvailabilityZoneState, zone string) ec2types.AvailabilityZone {
  	return createAvailabilityZoneWithID(region, state, zone, "id-"+zone)
  }
  
  func createAvailabilityZoneWithID(region string, state ec2types.AvailabilityZoneState, zone, zoneID string) ec2types.AvailabilityZone {
  	return ec2types.AvailabilityZone{
  		RegionName: aws.String(region),
  		State:      state,
  		ZoneName:   aws.String(zone),
  		ZoneId:     aws.String(zoneID),
  	}
  }
  