add unit test

Signed-off-by: Eytan Avisror <[email protected]>

Author: Eytan Avisror

controllers/provisioners/eks/eks_test.go

@@ -309,13 +309,29 @@ func MockTagDescription(key, value string) *autoscaling.TagDescription {
 	}
 }
 
+func MockTemplateOverrides(weight string, types ...string) []*autoscaling.LaunchTemplateOverrides {
+	overrides := make([]*autoscaling.LaunchTemplateOverrides, 0)
+	for _, t := range types {
+		overrides = append(overrides, &autoscaling.LaunchTemplateOverrides{
+			InstanceType:     aws.String(t),
+			WeightedCapacity: aws.String(weight),
+		})
+	}
+	return overrides
+}
+
 func MockScalingGroup(name string, withTemplate bool, t ...*autoscaling.TagDescription) *autoscaling.Group {
 	asg := &autoscaling.Group{
 		AutoScalingGroupName: aws.String(name),
 		Tags:                 t,
 		MinSize:              aws.Int64(3),
 		MaxSize:              aws.Int64(6),
 		VPCZoneIdentifier:    aws.String("subnet-1,subnet-2,subnet-3"),
+		Instances: []*autoscaling.Instance{
+			{
+				InstanceType: aws.String("m5.xlarge"),
+			},
+		},
 	}
 
 	if withTemplate {

controllers/provisioners/eks/helpers_test.go

@@ -806,6 +806,85 @@ func TestGetMountOpts(t *testing.T) {
 	}
 }
 
+func TestGetOverrides(t *testing.T) {
+	var (
+		g             = gomega.NewGomegaWithT(t)
+		k             = MockKubernetesClientSet()
+		ig            = MockInstanceGroup()
+		configuration = ig.GetEKSConfiguration()
+		asgMock       = NewAutoScalingMocker()
+		iamMock       = NewIamMocker()
+		eksMock       = NewEksMocker()
+		ec2Mock       = NewEc2Mocker()
+		ssmMock       = NewSsmMocker()
+	)
+
+	w := MockAwsWorker(asgMock, iamMock, eksMock, ec2Mock, ssmMock)
+	ctx := MockContext(ig, k, w)
+	state := ctx.GetDiscoveredState()
+
+	instancePool := "SubFamilyFlexible"
+
+	tests := []struct {
+		primaryType        string
+		scalingGroup       *autoscaling.Group
+		mixedInstancesSpec *v1alpha1.MixedInstancesPolicySpec
+		expectedOverrides  []*autoscaling.LaunchTemplateOverrides
+	}{
+		{
+			primaryType:  "m5.xlarge",
+			scalingGroup: MockScalingGroup("asg-1", true),
+			mixedInstancesSpec: &v1alpha1.MixedInstancesPolicySpec{
+				InstanceTypes: []*v1alpha1.InstanceTypeSpec{
+					{
+						Type:   "m5a.xlarge",
+						Weight: 1,
+					},
+					{
+						Type:   "m5g.xlarge",
+						Weight: 1,
+					},
+				},
+			},
+			expectedOverrides: MockTemplateOverrides("1", "m5a.xlarge", "m5g.xlarge", "m5.xlarge"),
+		},
+		{
+			primaryType:  "m5.xlarge",
+			scalingGroup: MockScalingGroup("asg-1", true),
+			mixedInstancesSpec: &v1alpha1.MixedInstancesPolicySpec{
+				InstancePool: &instancePool,
+			},
+			expectedOverrides: MockTemplateOverrides("1", "m5.xlarge"),
+		},
+		{
+			primaryType:  "t2.xlarge",
+			scalingGroup: MockScalingGroup("asg-1", true),
+			mixedInstancesSpec: &v1alpha1.MixedInstancesPolicySpec{
+				InstanceTypes: []*v1alpha1.InstanceTypeSpec{
+					{
+						Type:   "t2a.xlarge",
+						Weight: 1,
+					},
+					{
+						Type:   "t2g.xlarge",
+						Weight: 1,
+					},
+				},
+			},
+			expectedOverrides: MockTemplateOverrides("1", "t2a.xlarge", "t2g.xlarge", "t2.xlarge", "m5.xlarge"),
+		},
+	}
+
+	for i, tc := range tests {
+		t.Logf("Test #%v - %+v", i, tc.expectedOverrides)
+		configuration.MixedInstancesPolicy = tc.mixedInstancesSpec
+		state.ScalingGroup = tc.scalingGroup
+		ig.Spec.EKSSpec.EKSConfiguration.InstanceType = tc.primaryType
+		overrides := ctx.GetOverrides()
+		g.Expect(overrides).To(gomega.ConsistOf(tc.expectedOverrides))
+	}
+}
+
 func TestGetUserDataStages(t *testing.T) {
 	var (
 		g             = gomega.NewGomegaWithT(t)