Add test for building aws asg instance type with launch templates

This commit is contained in:
Sheldon Kwok 2018-07-10 13:41:49 -07:00
parent 8310086bde
commit b942eec809
3 changed files with 69 additions and 20 deletions

View File

@ -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),

View File

@ -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")

View File

@ -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