From 42cab3522475ee3a39f257d1ecee91dc0a61f18a Mon Sep 17 00:00:00 2001
From: Yu Xiang Zhang <yuxz@amazon.com>
Date: Thu, 14 Mar 2024 01:03:05 +0000
Subject: [PATCH] Include MixedInstancesPolicy LaunchTemplate for validation

---
 pkg/actions/nodegroup/scale.go      | 19 ++++++--
 pkg/actions/nodegroup/scale_test.go | 69 +++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+), 4 deletions(-)

diff --git a/pkg/actions/nodegroup/scale.go b/pkg/actions/nodegroup/scale.go
index 3faed6ce48..c1b97710d3 100644
--- a/pkg/actions/nodegroup/scale.go
+++ b/pkg/actions/nodegroup/scale.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/autoscaling"
+	autoscalingtypes "github.com/aws/aws-sdk-go-v2/service/autoscaling/types"
 	"github.com/aws/aws-sdk-go-v2/service/ec2"
 	awseks "github.com/aws/aws-sdk-go-v2/service/eks"
 	ekstypes "github.com/aws/aws-sdk-go-v2/service/eks/types"
@@ -162,10 +163,20 @@ func validateNodeGroupAMI(ctx context.Context, awsProvider api.ClusterProvider,
 	if len(asg.AutoScalingGroups) != 1 {
 		return fmt.Errorf("expected to find exactly one Auto Scaling group for nodegroup; got %d", len(asg.AutoScalingGroups))
 	}
-	lt := asg.AutoScalingGroups[0].LaunchTemplate
-	if lt == nil {
-		logger.Warning("nodegroup with Auto Scaling group %q does not have a launch template", asgName)
-		return nil
+
+	var lt *autoscalingtypes.LaunchTemplateSpecification
+	if asg.AutoScalingGroups[0].MixedInstancesPolicy == nil {
+		lt = asg.AutoScalingGroups[0].LaunchTemplate
+		if lt == nil {
+			logger.Warning("nodegroup with Auto Scaling group %q does not have a launch template", asgName)
+			return nil
+		}
+	} else {
+		// ASG only use the LT in MixedInstancesPolicy if the ASG has a MixedInstancesPolicy
+		if asg.AutoScalingGroups[0].MixedInstancesPolicy.LaunchTemplate == nil {
+			return fmt.Errorf("expected the MixedInstancesPolicy in Auto Scaling group %q to include a launch template", asgName)
+		}
+		lt = asg.AutoScalingGroups[0].MixedInstancesPolicy.LaunchTemplate.LaunchTemplateSpecification
 	}
 
 	ltData, err := awsProvider.EC2().DescribeLaunchTemplateVersions(ctx, &ec2.DescribeLaunchTemplateVersionsInput{
diff --git a/pkg/actions/nodegroup/scale_test.go b/pkg/actions/nodegroup/scale_test.go
index 0745451b6a..234725b19e 100644
--- a/pkg/actions/nodegroup/scale_test.go
+++ b/pkg/actions/nodegroup/scale_test.go
@@ -198,6 +198,75 @@ var _ = Describe("Scale", func() {
 			}
 		}
 
+		mockMixedInstanceNodeGroupAMI := func(missingLaunchTemplate bool, asgName string) {
+			asgOutputWithMixedInstancesPolicy := &autoscaling.DescribeAutoScalingGroupsOutput{
+				AutoScalingGroups: []autoscalingtypes.AutoScalingGroup{
+					{
+						MixedInstancesPolicy: &autoscalingtypes.MixedInstancesPolicy{},
+					},
+				},
+			}
+			if !missingLaunchTemplate {
+				asgOutputWithMixedInstancesPolicy.AutoScalingGroups[0].MixedInstancesPolicy.LaunchTemplate = &autoscalingtypes.LaunchTemplate{
+					LaunchTemplateSpecification: &autoscalingtypes.LaunchTemplateSpecification{
+						LaunchTemplateId:   aws.String("lt-1234"),
+						LaunchTemplateName: aws.String("eksctl-test-ng"),
+						Version:            aws.String("1"),
+					},
+				}
+			}
+			p.MockASG().On("DescribeAutoScalingGroups", mock.Anything, mock.MatchedBy(func(input *autoscaling.DescribeAutoScalingGroupsInput) bool {
+				return len(input.AutoScalingGroupNames) == 1 && input.AutoScalingGroupNames[0] == asgName
+			})).Return(asgOutputWithMixedInstancesPolicy, nil)
+			p.MockEC2().On("DescribeLaunchTemplateVersions", mock.Anything, mock.MatchedBy(func(input *ec2.DescribeLaunchTemplateVersionsInput) bool {
+				return len(input.Versions) == 1 && input.LaunchTemplateId != nil && *input.LaunchTemplateId == "lt-1234"
+			})).Return(&ec2.DescribeLaunchTemplateVersionsOutput{
+				LaunchTemplateVersions: []ec2types.LaunchTemplateVersion{
+					{
+						LaunchTemplateData: &ec2types.ResponseLaunchTemplateData{
+							ImageId: aws.String("ami-1234"),
+						},
+					},
+				},
+			}, nil)
+			describeImagesOutput := &ec2.DescribeImagesOutput{
+				Images: []ec2types.Image{
+					{
+						ImageId: aws.String("ami-1234"),
+					},
+				},
+			}
+			p.MockEC2().On("DescribeImages", mock.Anything, mock.MatchedBy(func(input *ec2.DescribeImagesInput) bool {
+				return len(input.ImageIds) == 1 && input.ImageIds[0] == "ami-1234"
+			})).Return(describeImagesOutput, nil)
+
+			p.MockASG().On("UpdateAutoScalingGroup", mock.Anything, &autoscaling.UpdateAutoScalingGroupInput{
+				AutoScalingGroupName: aws.String(asgName),
+				MinSize:              aws.Int32(1),
+				DesiredCapacity:      aws.Int32(3),
+			}).Return(nil, nil)
+		}
+
+		When("the ASG exists with a MixedInstancesPolicy", func() {
+			asgName := "asg-name"
+
+			BeforeEach(func() {
+				mockNodeGroupStack(ngName, asgName)
+			})
+
+			It("fails for missing launch template", func() {
+				mockMixedInstanceNodeGroupAMI(true, asgName)
+				err := m.Scale(context.Background(), ng, false)
+				Expect(err).To(MatchError(fmt.Sprintf("failed to scale nodegroup %q for cluster %q, error: expected the MixedInstancesPolicy in Auto Scaling group %q to include a launch template", ng.Name, clusterName, asgName)))
+			})
+
+			It("scales the nodegroup", func() {
+				mockMixedInstanceNodeGroupAMI(false, asgName)
+				err := m.Scale(context.Background(), ng, false)
+				Expect(err).NotTo(HaveOccurred())
+			})
+		})
+
 		When("the ASG exists", func() {
 			BeforeEach(func() {
 				mockNodeGroupStack(ngName, "asg-name")