From a06570751e4b332229ba573d130987d27c678f08 Mon Sep 17 00:00:00 2001
From: Andrew Starr-Bochicchio <a.starr.b@gmail.com>
Date: Wed, 7 Feb 2024 18:47:26 -0500
Subject: [PATCH] builder: fix distributing images to additional regions

---
 builder/digitalocean/builder_acc_test.go | 54 +++++++++++++++++++++++-
 builder/digitalocean/step_snapshot.go    |  6 ++-
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/builder/digitalocean/builder_acc_test.go b/builder/digitalocean/builder_acc_test.go
index da6f1e2..f88fb1d 100644
--- a/builder/digitalocean/builder_acc_test.go
+++ b/builder/digitalocean/builder_acc_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"os/exec"
 	"testing"
 
 	"github.com/digitalocean/godo"
@@ -20,6 +21,14 @@ func TestBuilderAcc_basic(t *testing.T) {
 	acctest.TestPlugin(t, &acctest.PluginTestCase{
 		Name:     "test-digitalocean-builder-basic",
 		Template: fmt.Sprintf(testBuilderAccBasic, "ubuntu-20-04-x64"),
+		Check: func(buildCommand *exec.Cmd, logfile string) error {
+			if buildCommand.ProcessState != nil {
+				if buildCommand.ProcessState.ExitCode() != 0 {
+					return fmt.Errorf("Bad exit code. Logfile: %s", logfile)
+				}
+			}
+			return nil
+		},
 	})
 }
 
@@ -30,6 +39,32 @@ func TestBuilderAcc_imageId(t *testing.T) {
 	acctest.TestPlugin(t, &acctest.PluginTestCase{
 		Name:     "test-digitalocean-builder-imageID",
 		Template: makeTemplateWithImageId(t),
+		Check: func(buildCommand *exec.Cmd, logfile string) error {
+			if buildCommand.ProcessState != nil {
+				if buildCommand.ProcessState.ExitCode() != 0 {
+					return fmt.Errorf("Bad exit code. Logfile: %s", logfile)
+				}
+			}
+			return nil
+		},
+	})
+}
+
+func TestBuilderAcc_multiRegion(t *testing.T) {
+	if skip := testAccPreCheck(t); skip == true {
+		return
+	}
+	acctest.TestPlugin(t, &acctest.PluginTestCase{
+		Name:     "test-digitalocean-builder-multi-region",
+		Template: fmt.Sprintf(testBuilderAccMultiRegion, "ubuntu-20-04-x64"),
+		Check: func(buildCommand *exec.Cmd, logfile string) error {
+			if buildCommand.ProcessState != nil {
+				if buildCommand.ProcessState.ExitCode() != 0 {
+					return fmt.Errorf("Bad exit code. Logfile: %s", logfile)
+				}
+			}
+			return nil
+		},
 	})
 }
 
@@ -84,7 +119,8 @@ func makeTemplateWithImageId(t *testing.T) string {
 	return ""
 }
 
-const testBuilderAccBasic = `
+const (
+	testBuilderAccBasic = `
 {
 	"builders": [{
 		"type": "digitalocean",
@@ -97,3 +133,19 @@ const testBuilderAccBasic = `
 	}]
 }
 `
+
+	testBuilderAccMultiRegion = `
+{
+	"builders": [{
+		"type": "digitalocean",
+		"region": "nyc2",
+		"size": "s-1vcpu-1gb",
+		"image": "%v",
+		"ssh_username": "root",
+		"user_data": "",
+		"user_data_file": "",
+		"snapshot_regions": ["nyc2", "nyc3"]
+	}]
+}
+`
+)
diff --git a/builder/digitalocean/step_snapshot.go b/builder/digitalocean/step_snapshot.go
index 4db457f..372b2b4 100644
--- a/builder/digitalocean/step_snapshot.go
+++ b/builder/digitalocean/step_snapshot.go
@@ -87,6 +87,8 @@ func (s *stepSnapshot) Run(ctx context.Context, state multistep.StateBag) multis
 				"type":   "transfer",
 				"region": snapshotRegions[transfer],
 			}
+
+			ui.Say(fmt.Sprintf("Transferring snapshot (ID: %d) to %s", images[0].ID, snapshotRegions[transfer]))
 			imageTransfer, _, err := client.ImageActions.Transfer(context.TODO(), images[0].ID, transferRequest)
 			if err != nil {
 				err := fmt.Errorf("Error transferring snapshot: %s", err)
@@ -94,8 +96,8 @@ func (s *stepSnapshot) Run(ctx context.Context, state multistep.StateBag) multis
 				ui.Error(err.Error())
 				return multistep.ActionHalt
 			}
-			ui.Say(fmt.Sprintf("transferring Snapshot ID: %d", imageTransfer.ID))
-			if err := WaitForImageState(godo.ActionCompleted, imageTransfer.ID, action.ID,
+
+			if err := WaitForImageState(godo.ActionCompleted, images[0].ID, imageTransfer.ID,
 				client, 20*time.Minute); err != nil {
 				// If we get an error the first time, actually report it
 				err := fmt.Errorf("Error waiting for snapshot transfer: %s", err)