add tests for step_ami_region_copy

This commit is contained in:
Megan Marsh 2019-05-03 14:47:09 -07:00
parent ae3cf48f71
commit 2c352cc681
7 changed files with 219 additions and 25 deletions

View File

@ -284,6 +284,7 @@ func (b *Builder) Run(ctx context.Context, ui packer.Ui, hook packer.Hook) (pack
RegionKeyIds: b.config.AMIRegionKMSKeyIDs,
EncryptBootVolume: b.config.AMIEncryptBootVolume,
Name: b.config.AMIName,
OriginalRegion: *ec2conn.Config.Region,
},
&awscommon.StepModifyAMIAttributes{
Description: b.config.AMIDescription,

View File

@ -10,6 +10,7 @@ import (
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/request"
"github.com/aws/aws-sdk-go/service/ec2"
"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
"github.com/hashicorp/packer/helper/multistep"
)
@ -36,7 +37,7 @@ type StateChangeConf struct {
// Following are wrapper functions that use Packer's environment-variables to
// determing retry logic, then call the AWS SDK's built-in waiters.
func WaitUntilAMIAvailable(ctx aws.Context, conn *ec2.EC2, imageId string) error {
func WaitUntilAMIAvailable(ctx aws.Context, conn ec2iface.EC2API, imageId string) error {
imageInput := ec2.DescribeImagesInput{
ImageIds: []*string{&imageId},
}

View File

@ -7,6 +7,7 @@ import (
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ec2"
"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
"github.com/hashicorp/packer/helper/multistep"
"github.com/hashicorp/packer/packer"
)
@ -18,26 +19,28 @@ type StepAMIRegionCopy struct {
RegionKeyIds map[string]string
EncryptBootVolume *bool // nil means preserve
Name string
OriginalRegion string
toDelete string
getRegionConn func(*AccessConfig, string) (ec2iface.EC2API, error)
}
func (s *StepAMIRegionCopy) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
ec2conn := state.Get("ec2").(*ec2.EC2)
ui := state.Get("ui").(packer.Ui)
amis := state.Get("amis").(map[string]string)
snapshots := state.Get("snapshots").(map[string][]string)
ami := amis[*ec2conn.Config.Region]
if s.EncryptBootVolume != nil {
// encrypt_boot was set, we now have to copy the temporary
ami := amis[s.OriginalRegion]
if s.EncryptBootVolume != nil && *s.EncryptBootVolume {
// encrypt_boot is true, so we have to copy the temporary
// AMI with required encryption setting.
// temp image was created by stepCreateAMI.
s.Regions = append(s.Regions, *ec2conn.Config.Region)
s.Regions = append(s.Regions, s.OriginalRegion)
if s.RegionKeyIds == nil {
s.RegionKeyIds = make(map[string]string)
}
s.RegionKeyIds[*ec2conn.Config.Region] = s.AMIKmsKeyId
s.RegionKeyIds[s.OriginalRegion] = s.AMIKmsKeyId
s.toDelete = ami
}
@ -45,7 +48,7 @@ func (s *StepAMIRegionCopy) Run(ctx context.Context, state multistep.StateBag) m
return multistep.ActionContinue
}
ui.Say(fmt.Sprintf("Copying AMI (%s) to other regions...", ami))
ui.Say(fmt.Sprintf("Copying/Encrypting AMI (%s) to other regions...", ami))
var lock sync.Mutex
var wg sync.WaitGroup
@ -54,15 +57,20 @@ func (s *StepAMIRegionCopy) Run(ctx context.Context, state multistep.StateBag) m
wg.Add(len(s.Regions))
for _, region := range s.Regions {
if region == *ec2conn.Config.Region &&
(s.EncryptBootVolume == nil || *s.EncryptBootVolume == false) {
if region == s.OriginalRegion {
if s.EncryptBootVolume == nil || *s.EncryptBootVolume == false {
ui.Message(fmt.Sprintf(
"Avoiding copying AMI to duplicate region %s", region))
wg.Done()
continue
} else {
// encryption is true and we're in the original region
ui.Message(fmt.Sprintf("Creating encrypted copy in build region: %s", region))
}
} else {
// in non-build region
ui.Message(fmt.Sprintf("Copying to: %s", region))
}
if s.EncryptBootVolume != nil && *s.EncryptBootVolume {
regKeyID = s.RegionKeyIds[region]
@ -70,7 +78,7 @@ func (s *StepAMIRegionCopy) Run(ctx context.Context, state multistep.StateBag) m
go func(region string) {
defer wg.Done()
id, snapshotIds, err := amiRegionCopy(ctx, state, s.AccessConfig, s.Name, ami, region, *ec2conn.Config.Region, regKeyID, s.EncryptBootVolume)
id, snapshotIds, err := s.amiRegionCopy(ctx, state, s.AccessConfig, s.Name, ami, region, s.OriginalRegion, regKeyID, s.EncryptBootVolume)
lock.Lock()
defer lock.Unlock()
amis[region] = id
@ -100,6 +108,11 @@ func (s *StepAMIRegionCopy) Cleanup(state multistep.StateBag) {
ec2conn := state.Get("ec2").(*ec2.EC2)
ui := state.Get("ui").(packer.Ui)
// cleanup is only for encrypted copies.
if s.EncryptBootVolume == nil || !*s.EncryptBootVolume {
return
}
// Delete the unencrypted amis and snapshots
ui.Say("Deregistering the AMI and deleting unencrypted temporary " +
"AMIs and snapshots")
@ -148,22 +161,34 @@ func (s *StepAMIRegionCopy) Cleanup(state multistep.StateBag) {
}
}
// amiRegionCopy does a copy for the given AMI to the target region and
// returns the resulting ID and snapshot IDs, or error.
func amiRegionCopy(ctx context.Context, state multistep.StateBag, config *AccessConfig, name, imageId,
target, source, keyId string, encrypt *bool) (string, []string, error) {
snapshotIds := []string{}
func getRegionConn(config *AccessConfig, target string) (ec2iface.EC2API, error) {
// Connect to the region where the AMI will be copied to
session, err := config.Session()
if err != nil {
return "", snapshotIds, err
return nil, fmt.Errorf("Error getting region connection for copy: %s", err)
}
regionconn := ec2.New(session.Copy(&aws.Config{
Region: aws.String(target),
}))
return regionconn, nil
}
// amiRegionCopy does a copy for the given AMI to the target region and
// returns the resulting ID and snapshot IDs, or error.
func (s *StepAMIRegionCopy) amiRegionCopy(ctx context.Context, state multistep.StateBag, config *AccessConfig, name, imageId,
target, source, keyId string, encrypt *bool) (string, []string, error) {
snapshotIds := []string{}
if s.getRegionConn == nil {
s.getRegionConn = getRegionConn
}
regionconn, err := s.getRegionConn(config, target)
if err != nil {
return "", snapshotIds, err
}
resp, err := regionconn.CopyImage(&ec2.CopyImageInput{
SourceRegion: &source,
SourceImageId: &imageId,

View File

@ -0,0 +1,164 @@
package common
import (
"bytes"
"context"
"fmt"
"testing"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/request"
"github.com/aws/aws-sdk-go/service/ec2"
"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
"github.com/hashicorp/packer/helper/multistep"
"github.com/hashicorp/packer/packer"
)
func boolPointer(tf bool) *bool {
return &tf
}
// Define a mock struct to be used in unit tests for common aws steps.
type mockEC2Conn struct {
ec2iface.EC2API
Config *aws.Config
// Counters to figure out what code path was taken
copyImageCount int
describeImagesCount int
deregisterImageCount int
deleteSnapshotCount int
waitCount int
}
func (m *mockEC2Conn) CopyImage(copyInput *ec2.CopyImageInput) (*ec2.CopyImageOutput, error) {
m.copyImageCount++
copiedImage := fmt.Sprintf("%s-copied-%d", *copyInput.SourceImageId, m.copyImageCount)
output := &ec2.CopyImageOutput{
ImageId: &copiedImage,
}
return output, nil
}
// functions we have to create mock responses for in order for test to run
func (m *mockEC2Conn) DescribeImages(*ec2.DescribeImagesInput) (*ec2.DescribeImagesOutput, error) {
m.describeImagesCount++
output := &ec2.DescribeImagesOutput{
Images: []*ec2.Image{{}},
}
return output, nil
}
func (m *mockEC2Conn) DeregisterImage(*ec2.DeregisterImageInput) (*ec2.DeregisterImageOutput, error) {
m.deregisterImageCount++
output := &ec2.DeregisterImageOutput{}
return output, nil
}
func (m *mockEC2Conn) DeleteSnapshot(*ec2.DeleteSnapshotInput) (*ec2.DeleteSnapshotOutput, error) {
m.deleteSnapshotCount++
output := &ec2.DeleteSnapshotOutput{}
return output, nil
}
func (m *mockEC2Conn) WaitUntilImageAvailableWithContext(aws.Context, *ec2.DescribeImagesInput, ...request.WaiterOption) error {
m.waitCount++
return nil
}
func getMockConn(config *AccessConfig, target string) (ec2iface.EC2API, error) {
mockConn := &mockEC2Conn{
Config: aws.NewConfig(),
}
return mockConn, nil
}
// Create statebag for running test
func tState() multistep.StateBag {
state := new(multistep.BasicStateBag)
state.Put("ui", &packer.BasicUi{
Reader: new(bytes.Buffer),
Writer: new(bytes.Buffer),
})
state.Put("amis", map[string]string{"us-east-1": "ami-12345"})
state.Put("snapshots", map[string][]string{"us-east-1": {"snap-0012345"}})
conn, _ := getMockConn(&AccessConfig{}, "us-east-2")
state.Put("ec2", conn)
return state
}
func TestStepAmiRegionCopy_nil_encryption(t *testing.T) {
// create step
stepAMIRegionCopy := StepAMIRegionCopy{
AccessConfig: testAccessConfig(),
Regions: make([]string, 0),
AMIKmsKeyId: "",
RegionKeyIds: make(map[string]string),
EncryptBootVolume: nil,
Name: "fake-ami-name",
OriginalRegion: "us-east-1",
}
// mock out the region connection code
stepAMIRegionCopy.getRegionConn = getMockConn
state := tState()
stepAMIRegionCopy.Run(context.Background(), state)
if stepAMIRegionCopy.toDelete != "" {
t.Fatalf("Shouldn't delete original AMI if not encrypted")
}
if len(stepAMIRegionCopy.Regions) > 0 {
t.Fatalf("Shouldn't have added original ami to original region")
}
}
func TestStepAmiRegionCopy_false_encryption(t *testing.T) {
// create step
stepAMIRegionCopy := StepAMIRegionCopy{
AccessConfig: testAccessConfig(),
Regions: make([]string, 0),
AMIKmsKeyId: "",
RegionKeyIds: make(map[string]string),
EncryptBootVolume: boolPointer(false),
Name: "fake-ami-name",
OriginalRegion: "us-east-1",
}
// mock out the region connection code
stepAMIRegionCopy.getRegionConn = getMockConn
state := tState()
stepAMIRegionCopy.Run(context.Background(), state)
if stepAMIRegionCopy.toDelete != "" {
t.Fatalf("Shouldn't delete original AMI if not encrypted")
}
if len(stepAMIRegionCopy.Regions) > 0 {
t.Fatalf("Shouldn't have added original ami to Regions")
}
}
func TestStepAmiRegionCopy_true_encryption(t *testing.T) {
// create step
stepAMIRegionCopy := StepAMIRegionCopy{
AccessConfig: testAccessConfig(),
Regions: make([]string, 0),
AMIKmsKeyId: "",
RegionKeyIds: make(map[string]string),
EncryptBootVolume: boolPointer(true),
Name: "fake-ami-name",
OriginalRegion: "us-east-1",
}
// mock out the region connection code
stepAMIRegionCopy.getRegionConn = getMockConn
state := tState()
stepAMIRegionCopy.Run(context.Background(), state)
if stepAMIRegionCopy.toDelete == "" {
t.Fatalf("Should delete original AMI if encrypted=true")
}
if len(stepAMIRegionCopy.Regions) == 0 {
t.Fatalf("Should have added original ami to Regions")
}
}

View File

@ -227,6 +227,7 @@ func (b *Builder) Run(ctx context.Context, ui packer.Ui, hook packer.Hook) (pack
RegionKeyIds: b.config.AMIRegionKMSKeyIDs,
EncryptBootVolume: b.config.AMIEncryptBootVolume,
Name: b.config.AMIName,
OriginalRegion: *ec2conn.Config.Region,
},
&awscommon.StepModifyAMIAttributes{
Description: b.config.AMIDescription,

View File

@ -252,6 +252,7 @@ func (b *Builder) Run(ctx context.Context, ui packer.Ui, hook packer.Hook) (pack
RegionKeyIds: b.config.AMIRegionKMSKeyIDs,
EncryptBootVolume: b.config.AMIEncryptBootVolume,
Name: b.config.AMIName,
OriginalRegion: *ec2conn.Config.Region,
},
&awscommon.StepModifyAMIAttributes{
Description: b.config.AMIDescription,

View File

@ -303,6 +303,7 @@ func (b *Builder) Run(ctx context.Context, ui packer.Ui, hook packer.Hook) (pack
RegionKeyIds: b.config.AMIRegionKMSKeyIDs,
EncryptBootVolume: b.config.AMIEncryptBootVolume,
Name: b.config.AMIName,
OriginalRegion: *ec2conn.Config.Region,
},
&awscommon.StepModifyAMIAttributes{
Description: b.config.AMIDescription,