From b12598c1201525a7a96bd0be8fb58e64ecfbeda5 Mon Sep 17 00:00:00 2001
From: kong <vucong2409@gmail.com>
Date: Sat, 11 May 2024 10:22:49 +0700
Subject: [PATCH] move spot_allocation_strategy check to prepare step

---
 builder/common/run_config.go             |  9 +++++++++
 builder/common/run_config_test.go        | 10 ++++++++++
 builder/common/step_run_spot_instance.go |  3 +--
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/builder/common/run_config.go b/builder/common/run_config.go
index e4250fbe2..e5de06398 100644
--- a/builder/common/run_config.go
+++ b/builder/common/run_config.go
@@ -11,9 +11,11 @@ import (
 	"net"
 	"os"
 	"regexp"
+	"slices"
 	"strings"
 	"time"
 
+	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/hashicorp/packer-plugin-sdk/communicator"
 	"github.com/hashicorp/packer-plugin-sdk/template/config"
 	confighelper "github.com/hashicorp/packer-plugin-sdk/template/config"
@@ -849,6 +851,13 @@ func (c *RunConfig) Prepare(ctx *interpolate.Context) []error {
 		}
 	}
 
+	if c.SpotAllocationStrategy != "" {
+		if slices.Contains(ec2.SpotAllocationStrategy_Values(), c.SpotAllocationStrategy) {
+			errs = append(errs, fmt.Errorf(
+				"Unknown spot_allocation_strategy: %s", c.SpotAllocationStrategy))
+		}
+	}
+
 	if c.UserData != "" && c.UserDataFile != "" {
 		errs = append(errs, fmt.Errorf("Only one of user_data or user_data_file can be specified."))
 	} else if c.UserDataFile != "" {
diff --git a/builder/common/run_config_test.go b/builder/common/run_config_test.go
index f29f35e3c..507a6996a 100644
--- a/builder/common/run_config_test.go
+++ b/builder/common/run_config_test.go
@@ -461,3 +461,13 @@ func TestRunConfigPrepare_InvalidTenantForHost(t *testing.T) {
 		})
 	}
 }
+
+func TestRunConfigPrepare_EnableSpotInstanceBadSpotAllocationStrategy(t *testing.T) {
+	c := testConfig()
+	// There should be some error when Spot Allocation Strategy is invalid.
+	c.SpotAllocationStrategy = "very-expensive-one"
+	err := c.Prepare(nil)
+	if len(err) != 1 {
+		t.Fatalf("Should error if spot_allocation_strategy is invalid.")
+	}
+}
diff --git a/builder/common/step_run_spot_instance.go b/builder/common/step_run_spot_instance.go
index 00e595994..afe685c19 100644
--- a/builder/common/step_run_spot_instance.go
+++ b/builder/common/step_run_spot_instance.go
@@ -9,7 +9,6 @@ import (
 	"fmt"
 	"io/ioutil"
 	"log"
-	"slices"
 	"strings"
 	"time"
 
@@ -385,7 +384,7 @@ func (s *StepRunSpotInstance) Run(ctx context.Context, state multistep.StateBag)
 		Type: aws.String("instant"),
 	}
 
-	if slices.Contains(ec2.SpotAllocationStrategy_Values(), s.SpotAllocationStrategy) {
+	if s.SpotAllocationStrategy != "" {
 		createFleetInput.SpotOptions = &ec2.SpotOptionsRequest{
 			AllocationStrategy: aws.String(s.SpotAllocationStrategy),
 		}