From 12d43c98f725311a917db299b8dd49cc14899292 Mon Sep 17 00:00:00 2001
From: Megan Marsh <megan@hashicorp.com>
Date: Tue, 27 Jun 2017 16:12:22 -0700
Subject: [PATCH] test the deduplication code

---
 .../amazon/ebssurrogate/step_register_ami.go  | 27 ++++++++------
 .../ebssurrogate/step_register_ami_test.go    | 37 +++++++++++++++++++
 2 files changed, 53 insertions(+), 11 deletions(-)
 create mode 100644 builder/amazon/ebssurrogate/step_register_ami_test.go

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