github.com/darmach/terratest@v0.34.8-0.20210517103231-80931f95e3ff/modules/aws/asg.go (about)

     1  package aws
     2  
     3  import (
     4  	"fmt"
     5  	"time"
     6  
     7  	"github.com/aws/aws-sdk-go/aws"
     8  	"github.com/aws/aws-sdk-go/service/autoscaling"
     9  	"github.com/stretchr/testify/require"
    10  
    11  	"github.com/gruntwork-io/terratest/modules/logger"
    12  	"github.com/gruntwork-io/terratest/modules/retry"
    13  	"github.com/gruntwork-io/terratest/modules/testing"
    14  )
    15  
    16  type AsgCapacityInfo struct {
    17  	MinCapacity     int64
    18  	MaxCapacity     int64
    19  	CurrentCapacity int64
    20  	DesiredCapacity int64
    21  }
    22  
    23  // GetCapacityInfoForAsg returns the capacity info for the queried asg as a struct, AsgCapacityInfo.
    24  func GetCapacityInfoForAsg(t testing.TestingT, asgName string, awsRegion string) AsgCapacityInfo {
    25  	capacityInfo, err := GetCapacityInfoForAsgE(t, asgName, awsRegion)
    26  	require.NoError(t, err)
    27  	return capacityInfo
    28  }
    29  
    30  // GetCapacityInfoForAsgE returns the capacity info for the queried asg as a struct, AsgCapacityInfo.
    31  func GetCapacityInfoForAsgE(t testing.TestingT, asgName string, awsRegion string) (AsgCapacityInfo, error) {
    32  	asgClient, err := NewAsgClientE(t, awsRegion)
    33  	if err != nil {
    34  		return AsgCapacityInfo{}, err
    35  	}
    36  
    37  	input := autoscaling.DescribeAutoScalingGroupsInput{AutoScalingGroupNames: []*string{aws.String(asgName)}}
    38  	output, err := asgClient.DescribeAutoScalingGroups(&input)
    39  	if err != nil {
    40  		return AsgCapacityInfo{}, err
    41  	}
    42  	groups := output.AutoScalingGroups
    43  	if len(groups) == 0 {
    44  		return AsgCapacityInfo{}, NewNotFoundError("ASG", asgName, awsRegion)
    45  	}
    46  	capacityInfo := AsgCapacityInfo{
    47  		MinCapacity:     *groups[0].MinSize,
    48  		MaxCapacity:     *groups[0].MaxSize,
    49  		DesiredCapacity: *groups[0].DesiredCapacity,
    50  		CurrentCapacity: int64(len(groups[0].Instances)),
    51  	}
    52  	return capacityInfo, nil
    53  }
    54  
    55  // GetInstanceIdsForAsg gets the IDs of EC2 Instances in the given ASG.
    56  func GetInstanceIdsForAsg(t testing.TestingT, asgName string, awsRegion string) []string {
    57  	ids, err := GetInstanceIdsForAsgE(t, asgName, awsRegion)
    58  	if err != nil {
    59  		t.Fatal(err)
    60  	}
    61  	return ids
    62  }
    63  
    64  // GetInstanceIdsForAsgE gets the IDs of EC2 Instances in the given ASG.
    65  func GetInstanceIdsForAsgE(t testing.TestingT, asgName string, awsRegion string) ([]string, error) {
    66  	asgClient, err := NewAsgClientE(t, awsRegion)
    67  	if err != nil {
    68  		return nil, err
    69  	}
    70  
    71  	input := autoscaling.DescribeAutoScalingGroupsInput{AutoScalingGroupNames: []*string{aws.String(asgName)}}
    72  	output, err := asgClient.DescribeAutoScalingGroups(&input)
    73  	if err != nil {
    74  		return nil, err
    75  	}
    76  
    77  	instanceIDs := []string{}
    78  	for _, asg := range output.AutoScalingGroups {
    79  		for _, instance := range asg.Instances {
    80  			instanceIDs = append(instanceIDs, aws.StringValue(instance.InstanceId))
    81  		}
    82  	}
    83  
    84  	return instanceIDs, nil
    85  }
    86  
    87  // WaitForCapacity waits for the currently set desired capacity to be reached on the ASG
    88  func WaitForCapacity(
    89  	t testing.TestingT,
    90  	asgName string,
    91  	region string,
    92  	maxRetries int,
    93  	sleepBetweenRetries time.Duration,
    94  ) {
    95  	err := WaitForCapacityE(t, asgName, region, maxRetries, sleepBetweenRetries)
    96  	require.NoError(t, err)
    97  }
    98  
    99  // WaitForCapacityE waits for the currently set desired capacity to be reached on the ASG
   100  func WaitForCapacityE(
   101  	t testing.TestingT,
   102  	asgName string,
   103  	region string,
   104  	maxRetries int,
   105  	sleepBetweenRetries time.Duration,
   106  ) error {
   107  	msg, err := retry.DoWithRetryE(
   108  		t,
   109  		fmt.Sprintf("Waiting for ASG %s to reach desired capacity.", asgName),
   110  		maxRetries,
   111  		sleepBetweenRetries,
   112  		func() (string, error) {
   113  			capacityInfo, err := GetCapacityInfoForAsgE(t, asgName, region)
   114  			if err != nil {
   115  				return "", err
   116  			}
   117  			if capacityInfo.CurrentCapacity != capacityInfo.DesiredCapacity {
   118  				return "", NewAsgCapacityNotMetError(asgName, capacityInfo.DesiredCapacity, capacityInfo.CurrentCapacity)
   119  			}
   120  			return fmt.Sprintf("ASG %s is now at desired capacity %d", asgName, capacityInfo.DesiredCapacity), nil
   121  		},
   122  	)
   123  	logger.Log(t, msg)
   124  	return err
   125  }
   126  
   127  // NewAsgClient creates an Auto Scaling Group client.
   128  func NewAsgClient(t testing.TestingT, region string) *autoscaling.AutoScaling {
   129  	client, err := NewAsgClientE(t, region)
   130  	if err != nil {
   131  		t.Fatal(err)
   132  	}
   133  	return client
   134  }
   135  
   136  // NewAsgClientE creates an Auto Scaling Group client.
   137  func NewAsgClientE(t testing.TestingT, region string) (*autoscaling.AutoScaling, error) {
   138  	sess, err := NewAuthenticatedSession(region)
   139  	if err != nil {
   140  		return nil, err
   141  	}
   142  
   143  	return autoscaling.New(sess), nil
   144  }