Specify aws in region request when creating a deployment (#1403)

* specify aws in region request * add test
astronomer · Oct 2, 2023 · ddfac75 · ddfac75
1 parent 323c832
commit ddfac75
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 0 deletions.
diff --git 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
@@ -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{