From 84f0079c3a2105c36f4e20ebb17f92e48f0b5bfb Mon Sep 17 00:00:00 2001
From: Andrew Starr-Bochicchio <a.starr.b@gmail.com>
Date: Wed, 13 May 2020 13:29:22 -0400
Subject: [PATCH] Move to a new func to make testable.

---
 builder/digitalocean/step_create_droplet.go   | 18 ++++++++-----
 .../digitalocean/step_create_droplet_test.go  | 26 +++++++++++++++++++
 2 files changed, 38 insertions(+), 6 deletions(-)
 create mode 100644 builder/digitalocean/step_create_droplet_test.go

diff --git a/builder/digitalocean/step_create_droplet.go b/builder/digitalocean/step_create_droplet.go
index e62a9ecb8..e246abf5e 100644
--- a/builder/digitalocean/step_create_droplet.go
+++ b/builder/digitalocean/step_create_droplet.go
@@ -37,12 +37,7 @@ func (s *stepCreateDroplet) Run(ctx context.Context, state multistep.StateBag) m
 		userData = string(contents)
 	}
 
-	createImage := godo.DropletCreateImage{Slug: c.Image}
-
-	imageId, err := strconv.Atoi(c.Image)
-	if err == nil {
-		createImage = godo.DropletCreateImage{ID: imageId}
-	}
+	createImage := getImageType(c.Image)
 
 	dropletCreateReq := &godo.DropletCreateRequest{
 		Name:   c.DropletName,
@@ -98,3 +93,14 @@ func (s *stepCreateDroplet) Cleanup(state multistep.StateBag) {
 			"Error destroying droplet. Please destroy it manually: %s", err))
 	}
 }
+
+func getImageType(image string) godo.DropletCreateImage {
+	createImage := godo.DropletCreateImage{Slug: image}
+
+	imageId, err := strconv.Atoi(image)
+	if err == nil {
+		createImage = godo.DropletCreateImage{ID: imageId}
+	}
+
+	return createImage
+}
diff --git a/builder/digitalocean/step_create_droplet_test.go b/builder/digitalocean/step_create_droplet_test.go
new file mode 100644
index 000000000..507d77d61
--- /dev/null
+++ b/builder/digitalocean/step_create_droplet_test.go
@@ -0,0 +1,26 @@
+package digitalocean
+
+import (
+	"testing"
+
+	"github.com/digitalocean/godo"
+)
+
+func TestBuilder_GetImageType(t *testing.T) {
+	imageTypeTests := []struct {
+		in  string
+		out godo.DropletCreateImage
+	}{
+		{"ubuntu-20-04-x64", godo.DropletCreateImage{Slug: "ubuntu-20-04-x64"}},
+		{"123456", godo.DropletCreateImage{ID: 123456}},
+	}
+
+	for _, tt := range imageTypeTests {
+		t.Run(tt.in, func(t *testing.T) {
+			i := getImageType(tt.in)
+			if i != tt.out {
+				t.Errorf("got %q, want %q", godo.Stringify(i), godo.Stringify(tt.out))
+			}
+		})
+	}
+}