allow pr preview logins

- `auth.Login()` succeeds with a valid pr preview env
- `auth.formatDomain()` and `auth.ValidateDomain()` allow pr preview
- `config.IsCloudContext()` allows pr preview
- add a pr prveiew context to `testing.NewTestConfig()`
- TODO - address a few TODOs

cloud/auth/auth.go

@@ -435,8 +435,7 @@ func Logout(domain string, out io.Writer) {
 
 func formatDomain(domain string) string {
 	if strings.Contains(domain, "cloud") {
-		splitDomain := strings.SplitN(domain, ".", splitNum) // This splits out 'cloud' from the domain string
-		domain = splitDomain[1]
+		domain = strings.Replace(domain, "cloud.", "", 1)
 	} else if domain == "" {
 		domain = Domain
 	}
@@ -445,24 +444,29 @@ func formatDomain(domain string) string {
 }
 
 func ValidateDomain(domain string) (astro.AuthConfig, error) {
-	validDomainsList := []string{"astronomer-dev.io", "astronomer-stage.io", "astronomer-perf.io", "astronomer.io", "localhost"}
-	var authConfig astro.AuthConfig
-
-	validDomain := false
-	for _, x := range validDomainsList {
-		if x == domain {
-			validDomain = true
-			break
-		}
+	var (
+		authConfig  astro.AuthConfig
+		prNum, addr string
+		validDomain bool
+	)
+
+	if strings.Contains(domain, "pr") {
+		prNum, domain, _ = strings.Cut(domain, ".")
 	}
+	validDomain = context.IsCloudDomain(domain)
 	if !validDomain {
 		return authConfig, errors.New("Error! Invalid domain. You are attempting to login into Astro. " +
 			"Are you trying to authenticate to Astronomer Software? If so, please change your current context with 'astro context switch'")
 	}
 
-	var addr string
-	if domain == "localhost" {
+	// TODO refactor and dry this out to remove the nolint
+	if domain == "localhost" { // nolint
 		addr = fmt.Sprintf("http://%s:8871/auth-config", domain)
+	} else if prNum != "" {
+		addr = fmt.Sprintf(
+			"https://%s.api.%s/hub/auth-config",
+			prNum, domain,
+		)
 	} else {
 		addr = fmt.Sprintf(
 			"https://api.%s/hub/auth-config",

cloud/auth/auth_test.go

@@ -87,6 +87,16 @@ func Test_validateDomain(t *testing.T) {
 	assert.Error(t, err)
 	assert.Errorf(t, err, "Error! Invalid domain. "+
 		"Are you trying to authenticate to Astronomer Software? If so, change your current context with 'astro context switch'. ")
+
+	t.Run("pr preview is a valid domain", func(t *testing.T) {
+		// TODO need to mock this as once this PR closes, test will fail
+		domain = "pr5723.astronomer-dev.io"
+		actual, err = ValidateDomain(domain)
+		assert.NoError(t, err)
+		assert.Equal(t, actual.ClientID, "GWCArPCS6pgFZbhnwx3WCDaJLBcvdFaV")
+		assert.Equal(t, actual.Audience, "astronomer-ee")
+		assert.Equal(t, actual.DomainURL, "https://astronomer-sandbox.us.auth0.com")
+	})
 }
 
 func TestOrgLookup(t *testing.T) {
@@ -455,6 +465,30 @@ func TestLogin(t *testing.T) {
 		err := Login("astronomer.io", "", "", mockClient, os.Stdout, false)
 		assert.NoError(t, err)
 	})
+	t.Run("can login to a pr preview environment successfully", func(t *testing.T) {
+		testUtil.InitTestConfig(testUtil.CloudPrPreview)
+		mockResponse := Result{RefreshToken: "test-token", AccessToken: "test-token", ExpiresIn: 300}
+		orgChecker := func(domain string) (string, error) {
+			return "test-org-id", nil
+		}
+		callbackHandler := func() (string, error) {
+			return "test-code", nil
+		}
+		tokenRequester := func(authConfig astro.AuthConfig, verifier, code string) (Result, error) {
+			return mockResponse, nil
+		}
+		openURL = func(url string) error {
+			return nil
+		}
+		authenticator = Authenticator{orgChecker, tokenRequester, callbackHandler}
+
+		mockClient := new(astro_mocks.Client)
+		mockClient.On("GetUserInfo").Return(mockSelfResp, nil).Once()
+		mockClient.On("ListWorkspaces", "test-org-id").Return([]astro.Workspace{{ID: "test-id"}}, nil).Once()
+
+		err := Login("pr5723.cloud.astronomer-dev.io", "", "", mockClient, os.Stdout, false)
+		assert.NoError(t, err)
+	})
 
 	t.Run("oauth token success", func(t *testing.T) {
 		mockClient := new(astro_mocks.Client)
@@ -686,3 +720,30 @@ func Test_writeResultToContext(t *testing.T) {
 	// test after changes
 	assertConfigContents("Bearer new_token", "new_refresh_token", time.Now().Add(1234*time.Second), "test.user@astronomer.io")
 }
+
+func TestFormatDomain(t *testing.T) {
+	t.Run("removes cloud from cloud.astronomer.io", func(t *testing.T) {
+		actual := formatDomain("cloud.astronomer.io")
+		assert.Equal(t, "astronomer.io", actual)
+	})
+	t.Run("removes cloud from cloud.astronomer-dev.io", func(t *testing.T) {
+		actual := formatDomain("cloud.astronomer-dev.io")
+		assert.Equal(t, "astronomer-dev.io", actual)
+	})
+	t.Run("removes cloud from cloud.astronomer-stage.io", func(t *testing.T) {
+		actual := formatDomain("cloud.astronomer-stage.io")
+		assert.Equal(t, "astronomer-stage.io", actual)
+	})
+	t.Run("removes cloud from cloud.astronomer-perf.io", func(t *testing.T) {
+		actual := formatDomain("cloud.astronomer-perf.io")
+		assert.Equal(t, "astronomer-perf.io", actual)
+	})
+	t.Run("removes cloud from pr1234.cloud.astronomer-dev.io", func(t *testing.T) {
+		actual := formatDomain("pr1234.cloud.astronomer-dev.io")
+		assert.Equal(t, "pr1234.astronomer-dev.io", actual)
+	})
+	t.Run("sets default domain if one was not provided", func(t *testing.T) {
+		actual := formatDomain("")
+		assert.Equal(t, "astronomer.io", actual)
+	})
+}

config/config.go

@@ -13,6 +13,7 @@ import (
 const (
 	CloudPlatform    = "cloud"
 	SoftwarePlatform = "software"
+	PrPreview        = "prprievew"
 
 	localhostDomain = "localhost"
 	astrohubDomain  = "astrohub"

context/context_test.go

@@ -60,6 +60,7 @@ func TestIsCloudContext(t *testing.T) {
 		{"software-domain", "dev.astrodev.com", config.SoftwarePlatform, false},
 		{"software-domain", "software.astronomer-test.io", config.SoftwarePlatform, false},
 		{"local-software", "localhost", config.SoftwarePlatform, false},
+		{"prpreview", "pr1234.cloud.astronomer-dev.io", config.PrPreview, true},
 	}
 
 	for _, tt := range tests {
@@ -112,3 +113,16 @@ func TestListContext(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Contains(t, buf.String(), "localhost")
 }
+
+func TestIsCloudDomain(t *testing.T) {
+	testUtil.InitTestConfig(testUtil.CloudPrPreview)
+	actual := IsCloudDomain("pr1234.cloud.astronomer.io")
+	assert.True(t, actual)
+
+	// TODO this makes the first find pass in viper.IsSet()
+	// SetContext("pr1234.cloud.astronomer.io")
+	// TODO this finds the context
+	Switch("pr1234.cloud.astronomer.io")
+	actual = IsCloudContext()
+	assert.True(t, actual)
+}

pkg/testing/testing.go

@@ -18,6 +18,7 @@ const (
 	CloudDevPlatform      = "dev"
 	CloudPerfPlatform     = "perf"
 	CloudStagePlatform    = "stage"
+	CloudPrPreview        = "prpreview"
 	SoftwarePlatform      = "software"
 	Initial               = "initial"
 	SQLCLI                = "sql_cli"
@@ -85,6 +86,8 @@ contexts:
 		testConfig = `beta:
   sql_cli: true
 `
+	case CloudPrPreview:
+		testConfig = fmt.Sprintf(testConfig, "pr1234.cloud.asrtronomer-dev.io", strings.Replace("pr1234.cloud.astronomer-dev.io", ".", "_", -1), "pr1234.cloud.asrtronomer-dev.io")
 	case ErrorReturningContext:
 		// this is an error returning case
 		testConfig = fmt.Sprintf(testConfig, "error", "error", "error")