test the deduplication code

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
+}

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