From ddfac75cc24944333186b901f77f563c302525cd Mon Sep 17 00:00:00 2001
From: David Koenitzer <davidkoenit@gmail.com>
Date: Mon, 2 Oct 2023 14:24:02 -0500
Subject: [PATCH] Specify aws in region request when creating a deployment
 (#1403)

* specify aws in region request

* add test
---
 cloud/deployment/deployment.go      |  4 +++
 cloud/deployment/deployment_test.go | 39 +++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/cloud/deployment/deployment.go b/cloud/deployment/deployment.go
index 5e38ee0d4..c9a9c684e 100644
--- a/cloud/deployment/deployment.go
+++ b/cloud/deployment/deployment.go
@@ -47,6 +47,7 @@ const (
 	CeleryExecutor = "CeleryExecutor"
 	notApplicable  = "N/A"
 	gcpCloud       = "gcp"
+	awsCloud       = "aws"
 	standard       = "standard"
 )
 
@@ -414,6 +415,9 @@ func ListClusterOptions(cloudProvider string, coreClient astrocore.CoreClient) (
 	if cloudProvider == gcpCloud {
 		provider = astrocore.GetClusterOptionsParamsProvider(astrocore.GetClusterOptionsParamsProviderGcp) //nolint
 	}
+	if cloudProvider == awsCloud {
+		provider = astrocore.GetClusterOptionsParamsProvider(astrocore.GetClusterOptionsParamsProviderAws) //nolint
+	}
 	optionsParams := &astrocore.GetClusterOptionsParams{
 		Provider: &provider,
 		Type:     astrocore.GetClusterOptionsParamsType(astrocore.GetClusterOptionsParamsTypeSHARED), //nolint
diff --git a/cloud/deployment/deployment_test.go b/cloud/deployment/deployment_test.go
index 211253182..06fef339b 100644
--- a/cloud/deployment/deployment_test.go
+++ b/cloud/deployment/deployment_test.go
@@ -545,6 +545,45 @@ func TestSelectRegion(t *testing.T) {
 		assert.Equal(t, region, resp)
 	})
 
+	t.Run("region via selection aws", func(t *testing.T) {
+		provider := astrocore.GetClusterOptionsParamsProvider(astrocore.GetClusterOptionsParamsProviderAws) //nolint
+		getSharedClusterOptionsParams := &astrocore.GetClusterOptionsParams{
+			Provider: &provider,
+			Type:     astrocore.GetClusterOptionsParamsType(astrocore.GetClusterOptionsParamsTypeSHARED), //nolint
+		}
+
+		mockOKRegionResponse := &astrocore.GetClusterOptionsResponse{
+			HTTPResponse: &http.Response{
+				StatusCode: 200,
+			},
+			JSON200: &[]astrocore.ClusterOptions{
+				{Regions: []astrocore.ProviderRegion{{Name: region}}},
+			},
+		}
+
+		mockCoreClient.On("GetClusterOptionsWithResponse", mock.Anything, getSharedClusterOptionsParams).Return(mockOKRegionResponse, nil).Once()
+
+		// mock os.Stdin
+		input := []byte("1")
+		r, w, err := os.Pipe()
+		if err != nil {
+			t.Fatal(err)
+		}
+		_, err = w.Write(input)
+		if err != nil {
+			t.Error(err)
+		}
+		w.Close()
+		stdin := os.Stdin
+		// Restore stdin right after the test.
+		defer func() { os.Stdin = stdin }()
+		os.Stdin = r
+
+		resp, err := selectRegion("aws", "", mockCoreClient)
+		assert.NoError(t, err)
+		assert.Equal(t, region, resp)
+	})
+
 	t.Run("region invalid selection", func(t *testing.T) {
 		provider := astrocore.GetClusterOptionsParamsProvider(astrocore.GetClusterOptionsParamsProviderGcp) //nolint
 		getSharedClusterOptionsParams := &astrocore.GetClusterOptionsParams{