amazon/chroot: Refactor step_register_ami for testing EBS-related behaviour

There is some logic around how EBS mappings are handled / generated depending
on whether an image if from scratch or not. There was no testing around
this behaviour before. This strips that logic out into a separate function,
to enable testing it independentlhy.

builder/amazon/chroot/step_register_ami.go

@@ -21,61 +21,19 @@ type StepRegisterAMI struct {
 func (s *StepRegisterAMI) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
 	config := state.Get("config").(*Config)
 	ec2conn := state.Get("ec2").(*ec2.EC2)
-	snapshotId := state.Get("snapshot_id").(string)
+	snapshotID := state.Get("snapshot_id").(string)
 	ui := state.Get("ui").(packer.Ui)
 
 	ui.Say("Registering the AMI...")
 
-	var (
-		registerOpts   *ec2.RegisterImageInput
-		mappings       []*ec2.BlockDeviceMapping
-		image          *ec2.Image
-		rootDeviceName string
-	)
+	var registerOpts *ec2.RegisterImageInput
 
+	// Source Image is only required to be passed if the image is not from scratch
 	if config.FromScratch {
-		mappings = config.AMIBlockDevices.BuildAMIDevices()
-		rootDeviceName = config.RootDeviceName
+		registerOpts = buildBaseRegisterOpts(config, nil, s.RootVolumeSize, snapshotID)
 	} else {
-		image = state.Get("source_image").(*ec2.Image)
-		mappings = image.BlockDeviceMappings
-		rootDeviceName = *image.RootDeviceName
-	}
-
-	newMappings := make([]*ec2.BlockDeviceMapping, len(mappings))
-	for i, device := range mappings {
-		newDevice := device
-		if *newDevice.DeviceName == rootDeviceName {
-			if newDevice.Ebs != nil {
-				newDevice.Ebs.SnapshotId = aws.String(snapshotId)
-			} else {
-				newDevice.Ebs = &ec2.EbsBlockDevice{SnapshotId: aws.String(snapshotId)}
-			}
-
-			if config.FromScratch || s.RootVolumeSize > *newDevice.Ebs.VolumeSize {
-				newDevice.Ebs.VolumeSize = aws.Int64(s.RootVolumeSize)
-			}
-		}
-
-		// assume working from a snapshot, so we unset the Encrypted field if set,
-		// otherwise AWS API will return InvalidParameter
-		if newDevice.Ebs != nil && newDevice.Ebs.Encrypted != nil {
-			newDevice.Ebs.Encrypted = nil
-		}
-
-		newMappings[i] = newDevice
-	}
-
-	if config.FromScratch {
-		registerOpts = &ec2.RegisterImageInput{
-			Name:                &config.AMIName,
-			Architecture:        aws.String(ec2.ArchitectureValuesX8664),
-			RootDeviceName:      aws.String(rootDeviceName),
-			VirtualizationType:  aws.String(config.AMIVirtType),
-			BlockDeviceMappings: newMappings,
-		}
-	} else {
-		registerOpts = buildRegisterOpts(config, image, newMappings)
+		image := state.Get("source_image").(*ec2.Image)
+		registerOpts = buildBaseRegisterOpts(config, image, s.RootVolumeSize, snapshotID)
 	}
 
 	if s.EnableAMISriovNetSupport {
@@ -115,6 +73,59 @@ func (s *StepRegisterAMI) Run(ctx context.Context, state multistep.StateBag) mul
 
 func (s *StepRegisterAMI) Cleanup(state multistep.StateBag) {}
 
+// Builds the base register opts with architecture, name, root block device, mappings, virtualizationtype
+func buildBaseRegisterOpts(config *Config, sourceImage *ec2.Image, rootVolumeSize int64, snapshotID string) *ec2.RegisterImageInput {
+	var (
+		mappings       []*ec2.BlockDeviceMapping
+		rootDeviceName string
+	)
+
+	if config.FromScratch {
+		mappings = config.AMIBlockDevices.BuildAMIDevices()
+		rootDeviceName = config.RootDeviceName
+	} else {
+		// If config.FromScratch is false, source image must be set
+		mappings = sourceImage.BlockDeviceMappings
+		rootDeviceName = *sourceImage.RootDeviceName
+	}
+
+	newMappings := make([]*ec2.BlockDeviceMapping, len(mappings))
+	for i, device := range mappings {
+		newDevice := device
+		if *newDevice.DeviceName == rootDeviceName {
+			if newDevice.Ebs != nil {
+				newDevice.Ebs.SnapshotId = aws.String(snapshotID)
+			} else {
+				newDevice.Ebs = &ec2.EbsBlockDevice{SnapshotId: aws.String(snapshotID)}
+			}
+
+			if config.FromScratch || rootVolumeSize > *newDevice.Ebs.VolumeSize {
+				newDevice.Ebs.VolumeSize = aws.Int64(rootVolumeSize)
+			}
+		}
+
+		// assume working from a snapshot, so we unset the Encrypted field if set,
+		// otherwise AWS API will return InvalidParameter
+		if newDevice.Ebs != nil && newDevice.Ebs.Encrypted != nil {
+			newDevice.Ebs.Encrypted = nil
+		}
+
+		newMappings[i] = newDevice
+	}
+
+	if config.FromScratch {
+		return &ec2.RegisterImageInput{
+			Name:                &config.AMIName,
+			Architecture:        aws.String(ec2.ArchitectureValuesX8664),
+			RootDeviceName:      aws.String(rootDeviceName),
+			VirtualizationType:  aws.String(config.AMIVirtType),
+			BlockDeviceMappings: newMappings,
+		}
+	}
+
+	return buildRegisterOpts(config, sourceImage, newMappings)
+}
+
 func buildRegisterOpts(config *Config, image *ec2.Image, mappings []*ec2.BlockDeviceMapping) *ec2.RegisterImageInput {
 	registerOpts := &ec2.RegisterImageInput{
 		Name:                &config.AMIName,

builder/amazon/chroot/step_register_ami_test.go

@@ -1,6 +1,8 @@
 package chroot
 
 import (
+	amazon "github.com/hashicorp/packer/builder/amazon/common"
+	"github.com/hashicorp/packer/common"
 	"testing"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -71,3 +73,76 @@ func TestStepRegisterAmi_buildRegisterOpts_hvm(t *testing.T) {
 		t.Fatalf("Unexpected KernelId value: expected nil got %s\n", *opts.KernelId)
 	}
 }
+
+func TestStepRegisterAmi_buildRegisterOptsFromScratch(t *testing.T) {
+	rootDeviceName := "/dev/sda"
+	snapshotID := "foo"
+	config := Config{
+		FromScratch:  true,
+		PackerConfig: common.PackerConfig{},
+		AMIBlockDevices: amazon.AMIBlockDevices{
+			AMIMappings: []amazon.BlockDevice{
+				{
+					DeviceName: rootDeviceName,
+				},
+			},
+		},
+		RootDeviceName: rootDeviceName,
+	}
+	registerOpts := buildBaseRegisterOpts(&config, nil, 10, snapshotID)
+
+	if len(registerOpts.BlockDeviceMappings) != 1 {
+		t.Fatal("Expected block device mapping of length 1")
+	}
+
+	if *registerOpts.BlockDeviceMappings[0].Ebs.SnapshotId != snapshotID {
+		t.Fatalf("Snapshot ID of root disk not set to snapshot id %s", rootDeviceName)
+	}
+}
+
+func TestStepRegisterAmi_buildRegisterOptFromExistingImage(t *testing.T) {
+	rootDeviceName := "/dev/sda"
+	snapshotID := "foo"
+
+	config := Config{
+		FromScratch:  false,
+		PackerConfig: common.PackerConfig{},
+	}
+	sourceImage := ec2.Image{
+		RootDeviceName: &rootDeviceName,
+		BlockDeviceMappings: []*ec2.BlockDeviceMapping{
+			{
+				DeviceName: &rootDeviceName,
+				Ebs: &ec2.EbsBlockDevice{
+					VolumeSize: aws.Int64(10),
+				},
+			},
+			// Throw in an ephemeral device, it seems like all devices in the return struct in a source AMI have
+			// a size, even if it's for ephemeral
+			{
+				DeviceName:  aws.String("/dev/sdb"),
+				VirtualName: aws.String("ephemeral0"),
+				Ebs: &ec2.EbsBlockDevice{
+					VolumeSize: aws.Int64(0),
+				},
+			},
+		},
+	}
+	registerOpts := buildBaseRegisterOpts(&config, &sourceImage, 15, snapshotID)
+
+	if len(registerOpts.BlockDeviceMappings) != 2 {
+		t.Fatal("Expected block device mapping of length 2")
+	}
+
+	for _, dev := range registerOpts.BlockDeviceMappings {
+		if dev.Ebs.SnapshotId != nil && *dev.Ebs.SnapshotId == snapshotID {
+			// Even though root volume size is in config, it isn't used, instead we use the root volume size
+			// that's derived when we build the step
+			if *dev.Ebs.VolumeSize != 15 {
+				t.Fatalf("Root volume size not 15 GB instead %d", *dev.Ebs.VolumeSize)
+			}
+			return
+		}
+	}
+	t.Fatalf("Could not find device with snapshot ID %s", snapshotID)
+}