Move to a new func to make testable.

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

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