diff --git a/cluster-autoscaler/cloudprovider/aws/aws_cloud_provider_test.go b/cluster-autoscaler/cloudprovider/aws/aws_cloud_provider_test.go index c939e1a4dd..51a2e491d5 100644 --- a/cluster-autoscaler/cloudprovider/aws/aws_cloud_provider_test.go +++ b/cluster-autoscaler/cloudprovider/aws/aws_cloud_provider_test.go @@ -21,6 +21,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/autoscaling" + "github.com/aws/aws-sdk-go/service/ec2" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" apiv1 "k8s.io/api/core/v1" @@ -56,6 +57,15 @@ func (a *AutoScalingMock) TerminateInstanceInAutoScalingGroup(input *autoscaling return args.Get(0).(*autoscaling.TerminateInstanceInAutoScalingGroupOutput), nil } +type EC2Mock struct { + mock.Mock +} + +func (e *EC2Mock) DescribeLaunchTemplateVersions(i *ec2.DescribeLaunchTemplateVersionsInput) (*ec2.DescribeLaunchTemplateVersionsOutput, error) { + args := e.Called(i) + return args.Get(0).(*ec2.DescribeLaunchTemplateVersionsOutput), nil +} + var testService = autoScalingWrapper{&AutoScalingMock{}} var testAwsManager = &AwsManager{ @@ -66,13 +76,13 @@ var testAwsManager = &AwsManager{ interrupt: make(chan struct{}), service: testService, }, - service: testService, + autoScalingService: testService, } func newTestAwsManagerWithService(service autoScaling, autoDiscoverySpecs []cloudprovider.ASGAutoDiscoveryConfig) *AwsManager { wrapper := autoScalingWrapper{service} return &AwsManager{ - service: wrapper, + autoScalingService: wrapper, asgCache: &asgCache{ registeredAsgs: make([]*asg, 0), asgToInstances: make(map[AwsRef][]AwsInstanceRef), diff --git a/cluster-autoscaler/cloudprovider/aws/aws_manager.go b/cluster-autoscaler/cloudprovider/aws/aws_manager.go index 465656020a..1bbfeed99e 100644 --- a/cluster-autoscaler/cloudprovider/aws/aws_manager.go +++ b/cluster-autoscaler/cloudprovider/aws/aws_manager.go @@ -50,10 +50,10 @@ const ( // AwsManager is handles aws communication and data caching. type AwsManager struct { - service autoScalingWrapper - ec2 ec2Wrapper - asgCache *asgCache - lastRefresh time.Time + autoScalingService autoScalingWrapper + ec2Service ec2Wrapper + asgCache *asgCache + lastRefresh time.Time } type asgTemplate struct { @@ -67,7 +67,8 @@ type asgTemplate struct { func createAWSManagerInternal( configReader io.Reader, discoveryOpts cloudprovider.NodeGroupDiscoveryOptions, - service *autoScalingWrapper, + autoScalingService *autoScalingWrapper, + ec2Service *ec2Wrapper, ) (*AwsManager, error) { if configReader != nil { var cfg provider_aws.CloudConfig @@ -77,9 +78,15 @@ func createAWSManagerInternal( } } - if service == nil { - service = &autoScalingWrapper{ - autoscaling.New(session.New()), + if autoScalingService == nil || ec2Service == nil { + sess := session.New() + + if autoScalingService == nil { + autoScalingService = &autoScalingWrapper{autoscaling.New(sess)} + } + + if ec2Service == nil { + ec2Service = &ec2Wrapper{ec2.New(sess)} } } @@ -88,15 +95,15 @@ func createAWSManagerInternal( return nil, err } - cache, err := newASGCache(*service, discoveryOpts.NodeGroupSpecs, specs) + cache, err := newASGCache(*autoScalingService, discoveryOpts.NodeGroupSpecs, specs) if err != nil { return nil, err } manager := &AwsManager{ - service: *service, - ec2: ec2Wrapper{ec2.New(session.New())}, - asgCache: cache, + autoScalingService: *autoScalingService, + ec2Service: *ec2Service, + asgCache: cache, } if err := manager.forceRefresh(); err != nil { @@ -108,7 +115,7 @@ func createAWSManagerInternal( // CreateAwsManager constructs awsManager object. func CreateAwsManager(configReader io.Reader, discoveryOpts cloudprovider.NodeGroupDiscoveryOptions) (*AwsManager, error) { - return createAWSManagerInternal(configReader, discoveryOpts, nil) + return createAWSManagerInternal(configReader, discoveryOpts, nil, nil) } // Refresh is called before every main loop and can be used to dynamically update cloud provider state. @@ -152,7 +159,7 @@ func (m *AwsManager) SetAsgSize(asg *asg, size int) error { HonorCooldown: aws.Bool(false), } glog.V(0).Infof("Setting asg %s size to %d", asg.Name, size) - _, err := m.service.SetDesiredCapacity(params) + _, err := m.autoScalingService.SetDesiredCapacity(params) if err != nil { return err } @@ -187,7 +194,7 @@ func (m *AwsManager) DeleteInstances(instances []*AwsInstanceRef) error { InstanceId: aws.String(instance.Name), ShouldDecrementDesiredCapacity: aws.Bool(true), } - resp, err := m.service.TerminateInstanceInAutoScalingGroup(params) + resp, err := m.autoScalingService.TerminateInstanceInAutoScalingGroup(params) if err != nil { return err } @@ -229,9 +236,9 @@ func (m *AwsManager) getAsgTemplate(asg *asg) (*asgTemplate, error) { func (m *AwsManager) buildInstanceType(asg *asg) (string, error) { if asg.LaunchConfigurationName != "" { - return m.service.getInstanceTypeByLCName(asg.LaunchConfigurationName) + return m.autoScalingService.getInstanceTypeByLCName(asg.LaunchConfigurationName) } else if asg.LaunchTemplateName != "" && asg.LaunchTemplateVersion != "" { - return m.ec2.getInstanceTypeByLT(asg.LaunchTemplateName, asg.LaunchTemplateVersion) + return m.ec2Service.getInstanceTypeByLT(asg.LaunchTemplateName, asg.LaunchTemplateVersion) } return "", fmt.Errorf("Unable to get instance type from launch config or launch template") diff --git a/cluster-autoscaler/cloudprovider/aws/aws_manager_test.go b/cluster-autoscaler/cloudprovider/aws/aws_manager_test.go index 2abcaca916..8eddd7eb53 100644 --- a/cluster-autoscaler/cloudprovider/aws/aws_manager_test.go +++ b/cluster-autoscaler/cloudprovider/aws/aws_manager_test.go @@ -22,6 +22,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/autoscaling" + "github.com/aws/aws-sdk-go/service/ec2" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" apiv1 "k8s.io/api/core/v1" @@ -159,7 +160,7 @@ func TestFetchExplicitAsgs(t *testing.T) { }, } // fetchExplicitASGs is called at manager creation time. - m, err := createAWSManagerInternal(nil, do, &autoScalingWrapper{s}) + m, err := createAWSManagerInternal(nil, do, &autoScalingWrapper{s}, nil) assert.NoError(t, err) asgs := m.asgCache.Get() @@ -167,6 +168,37 @@ func TestFetchExplicitAsgs(t *testing.T) { validateAsg(t, asgs[0], groupname, min, max) } +func TestBuildInstanceType(t *testing.T) { + ltName, ltVersion, instanceType := "launcher", "1", "t2.large" + + s := &EC2Mock{} + s.On("DescribeLaunchTemplateVersions", &ec2.DescribeLaunchTemplateVersionsInput{ + LaunchTemplateName: aws.String(ltName), + Versions: []*string{aws.String(ltVersion)}, + }).Return(&ec2.DescribeLaunchTemplateVersionsOutput{ + LaunchTemplateVersions: []*ec2.LaunchTemplateVersion{ + { + LaunchTemplateData: &ec2.ResponseLaunchTemplateData{ + InstanceType: aws.String(instanceType), + }, + }, + }, + }) + + m, err := createAWSManagerInternal(nil, cloudprovider.NodeGroupDiscoveryOptions{}, nil, &ec2Wrapper{s}) + assert.NoError(t, err) + + asg := asg{ + LaunchTemplateName: ltName, + LaunchTemplateVersion: ltVersion, + } + + builtInstanceType, err := m.buildInstanceType(&asg) + + assert.NoError(t, err) + assert.Equal(t, instanceType, builtInstanceType) +} + /* Disabled due to flakiness. See https://github.com/kubernetes/autoscaler/issues/608 func TestFetchAutoAsgs(t *testing.T) { min, max := 1, 10