add tests for step_ami_region_copy

builder/amazon/chroot/builder.go

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

builder/amazon/common/state.go

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

builder/amazon/common/step_ami_region_copy.go

@@ -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
+	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 {
+	ami := amis[s.OriginalRegion]
-		// encrypt_boot was set, we now have to copy the temporary
+
+	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,23 +57,28 @@ 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 &&
+		if region == s.OriginalRegion {
-			(s.EncryptBootVolume == nil || *s.EncryptBootVolume == false) {
+			if s.EncryptBootVolume == nil || *s.EncryptBootVolume == false {
-			ui.Message(fmt.Sprintf(
+				ui.Message(fmt.Sprintf(
-				"Avoiding copying AMI to duplicate region %s", region))
+					"Avoiding copying AMI to duplicate region %s", region))
-			wg.Done()
+				wg.Done()
-			continue
+				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))
 		}
 
-		ui.Message(fmt.Sprintf("Copying to: %s", region))
-
 		if s.EncryptBootVolume != nil && *s.EncryptBootVolume {
 			regKeyID = s.RegionKeyIds[region]
 		}
 
 		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
+func getRegionConn(config *AccessConfig, target string) (ec2iface.EC2API, error) {
-// 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{}
-
 	// 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,

builder/amazon/common/step_ami_region_copy_test.go

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

builder/amazon/ebs/builder.go

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

builder/amazon/ebssurrogate/builder.go

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

builder/amazon/instance/builder.go

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