diff --git a/builder/amazon/ebssurrogate/step_register_ami.go b/builder/amazon/ebssurrogate/step_register_ami.go index 0b9875308..d2f0f2250 100644 --- a/builder/amazon/ebssurrogate/step_register_ami.go +++ b/builder/amazon/ebssurrogate/step_register_ami.go @@ -25,17 +25,7 @@ func (s *StepRegisterAMI) Run(state multistep.StateBag) multistep.StepAction { ui.Say("Registering the AMI...") - // Defensive coding to make sure we only add the root volume once - blockDevicesExcludingRoot := make([]*ec2.BlockDeviceMapping, 0, len(s.BlockDevices)) - for _, blockDevice := range s.BlockDevices { - if *blockDevice.DeviceName == s.RootDevice.SourceDeviceName { - continue - } - - blockDevicesExcludingRoot = append(blockDevicesExcludingRoot, blockDevice) - } - - blockDevicesExcludingRoot = append(blockDevicesExcludingRoot, s.RootDevice.createBlockDeviceMapping(snapshotId)) + blockDevicesExcludingRoot := DeduplicateRootVolume(s.BlockDevices, s.RootDevice, snapshotId) registerOpts := &ec2.RegisterImageInput{ Name: &config.AMIName, @@ -126,3 +116,18 @@ func (s *StepRegisterAMI) Cleanup(state multistep.StateBag) { return } } + +func DeduplicateRootVolume(BlockDevices []*ec2.BlockDeviceMapping, RootDevice RootBlockDevice, snapshotId string) []*ec2.BlockDeviceMapping { + // Defensive coding to make sure we only add the root volume once + blockDevicesExcludingRoot := make([]*ec2.BlockDeviceMapping, 0, len(BlockDevices)) + for _, blockDevice := range BlockDevices { + if *blockDevice.DeviceName == RootDevice.SourceDeviceName { + continue + } + + blockDevicesExcludingRoot = append(blockDevicesExcludingRoot, blockDevice) + } + + blockDevicesExcludingRoot = append(blockDevicesExcludingRoot, RootDevice.createBlockDeviceMapping(snapshotId)) + return blockDevicesExcludingRoot +} diff --git a/builder/amazon/ebssurrogate/step_register_ami_test.go b/builder/amazon/ebssurrogate/step_register_ami_test.go new file mode 100644 index 000000000..b2abd7827 --- /dev/null +++ b/builder/amazon/ebssurrogate/step_register_ami_test.go @@ -0,0 +1,37 @@ +package ebssurrogate + +import ( + "testing" + + "github.com/aws/aws-sdk-go/service/ec2" +) + +func GetStringPointer() *string { + tmp := "/dev/name" + return &tmp +} + +func GetTestDevice() *ec2.BlockDeviceMapping { + TestDev := ec2.BlockDeviceMapping{ + DeviceName: GetStringPointer(), + } + return &TestDev +} + +func TestStepRegisterAmi_DeduplicateRootVolume(t *testing.T) { + TestRootDevice := RootBlockDevice{} + TestRootDevice.SourceDeviceName = "/dev/name" + + blockDevices := []*ec2.BlockDeviceMapping{} + blockDevicesExcludingRoot := DeduplicateRootVolume(blockDevices, TestRootDevice, "12342351") + if len(blockDevicesExcludingRoot) != 1 { + t.Fatalf("Unexpected length of block devices list") + } + + TestBlockDevice := GetTestDevice() + blockDevices = append(blockDevices, TestBlockDevice) + blockDevicesExcludingRoot = DeduplicateRootVolume(blockDevices, TestRootDevice, "12342351") + if len(blockDevicesExcludingRoot) != 1 { + t.Fatalf("Unexpected length of block devices list") + } +}