log into pr preview environments (#865)

* 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()`

* add domainUtil with domain and url helpers

- anytime we need to check domain or URLs to be derived from a domain for the user login flow, we use the `domainutil` package
- tighten up the domain and pr preview regexes
- able to transform astrohub url -> core api url for valid domains
- remove redundant `GetCloudAPIURL()`

* fix deployment dashboard url

- `deploy.go` and `deplyment.go` use `GetDeploymentURL()` to print the dashboard url
- resolves #806
- resolves #807
- resolves #805
- resolves #808

cloud/auth/auth.go

@@ -22,14 +22,15 @@ import (
 	"github.com/astronomer/astro-cli/config"
 	"github.com/astronomer/astro-cli/context"
 	"github.com/astronomer/astro-cli/pkg/ansi"
+	"github.com/astronomer/astro-cli/pkg/domainutil"
 	"github.com/astronomer/astro-cli/pkg/httputil"
 	"github.com/astronomer/astro-cli/pkg/input"
 	"github.com/astronomer/astro-cli/pkg/util"
 )
 
 const (
-	Domain      = "astronomer.io"
-	localDomain = "localhost"
+	authConfigEndpoint = "auth-config"
+	orgLookupEndpoint  = "organization-lookup"
 
 	cliChooseWorkspace     = "Please choose a workspace:"
 	cliSetWorkspaceExample = "\nNo default workspace detected, you can list workspaces with \n\tastro workspace list\nand set your default workspace with \n\tastro workspace switch [WORKSPACEID]\n\n"
@@ -47,7 +48,6 @@ var (
 
 var (
 	err             error
-	splitNum        = 2
 	callbackChannel = make(chan string, 1)
 	callbackTimeout = time.Second * 300
 	redirectURI     = "http://localhost:12345/callback"
@@ -61,20 +61,8 @@ var authenticator = Authenticator{
 }
 
 func orgLookup(domain string) (string, error) {
-	if strings.Contains(domain, "cloud") { // case when the domain has cloud as prefix, i.e. cloud.astronomer.io
-		splitDomain := strings.SplitN(domain, ".", splitNum)
-		domain = splitDomain[1]
-	}
-	var addr string
-	if domain == localDomain {
-		addr = "http://localhost:8871/organization-lookup"
-	} else {
-		addr = fmt.Sprintf(
-			"%s://api.%s/hub/organization-lookup",
-			config.CFG.CloudAPIProtocol.GetString(),
-			domain,
-		)
-	}
+	addr := domainutil.GetURLToEndpoint("https", domain, orgLookupEndpoint)
+
 	ctx := http_context.Background()
 	reqData, err := json.Marshal(orgLookupRequest{Email: userEmail})
 	if err != nil {
@@ -362,7 +350,7 @@ func checkToken(c *config.Context, client astro.Client, out io.Writer) error {
 // Login handles authentication to astronomer api and registry
 func Login(domain, orgID, token string, client astro.Client, out io.Writer, shouldDisplayLoginLink bool) error {
 	var res Result
-	domain = formatDomain(domain)
+	domain = domainutil.FormatDomain(domain)
 	authConfig, err := ValidateDomain(domain)
 	if err != nil {
 		return err
@@ -386,7 +374,6 @@ func Login(domain, orgID, token string, client astro.Client, out io.Writer, shou
 		}
 	}
 
-	// If no domain specified
 	// Create context if it does not exist
 	if domain != "" {
 		// Switch context now that we ensured context exists
@@ -437,42 +424,21 @@ func Logout(domain string, out io.Writer) {
 	fmt.Fprintln(out, "Successfully logged out of Astronomer")
 }
 
-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]
-	} else if domain == "" {
-		domain = Domain
-	}
-
-	return domain
-}
-
 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
+		addr        string
+		validDomain bool
+	)
+
+	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" {
-		addr = fmt.Sprintf("http://%s:8871/auth-config", domain)
-	} else {
-		addr = fmt.Sprintf(
-			"https://api.%s/hub/auth-config",
-			domain,
-		)
-	}
+	addr = domainutil.GetURLToEndpoint("https", domain, authConfigEndpoint)
+
 	ctx := http_context.Background()
 	doOptions := &httputil.DoOptions{
 		Context: ctx,

cloud/auth/auth_test.go

@@ -87,6 +87,30 @@ 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) {
+		// mocking this as once a PR closes, test would fail
+		mockResponse := astro.AuthConfig{
+			ClientID:  "client-id",
+			Audience:  "audience",
+			DomainURL: "https://myURL.com/",
+		}
+		jsonResponse, err := json.Marshal(mockResponse)
+		assert.NoError(t, err)
+		httpClient = testUtil.NewTestClient(func(req *http.Request) *http.Response {
+			return &http.Response{
+				StatusCode: 200,
+				Body:       io.NopCloser(bytes.NewBuffer(jsonResponse)),
+				Header:     make(http.Header),
+			}
+		})
+		domain = "pr1234.astronomer-dev.io"
+		actual, err = ValidateDomain(domain)
+		assert.NoError(t, err)
+		assert.Equal(t, actual.ClientID, mockResponse.ClientID)
+		assert.Equal(t, actual.Audience, mockResponse.Audience)
+		assert.Equal(t, actual.DomainURL, mockResponse.DomainURL)
+	})
 }
 
 func TestOrgLookup(t *testing.T) {
@@ -459,6 +483,45 @@ 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)
+		// mocking this as once a PR closes, test would fail
+		mockAuthConfigResponse := astro.AuthConfig{
+			ClientID:  "client-id",
+			Audience:  "audience",
+			DomainURL: "https://myURL.com/",
+		}
+		jsonResponse, err := json.Marshal(mockAuthConfigResponse)
+		assert.NoError(t, err)
+		httpClient = testUtil.NewTestClient(func(req *http.Request) *http.Response {
+			return &http.Response{
+				StatusCode: 200,
+				Body:       io.NopCloser(bytes.NewBuffer(jsonResponse)),
+				Header:     make(http.Header),
+			}
+		})
+		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)
@@ -521,8 +584,11 @@ func TestLogin(t *testing.T) {
 			Label: "something-label",
 		}}, nil).Once()
 		// initialize the test authenticator
+		orgChecker := func(domain string) (string, error) {
+			return "test-org-id", nil
+		}
 		authenticator = Authenticator{
-			orgChecker:      orgLookup,
+			orgChecker:      orgChecker,
 			callbackHandler: func() (string, error) { return "authorizationCode", nil },
 			tokenRequester: func(authConfig astro.AuthConfig, verifier, code string) (Result, error) {
 				return Result{
@@ -550,8 +616,11 @@ func TestLogin(t *testing.T) {
 			Label: "something-label",
 		}}, nil).Once()
 		// initialize the test authenticator
+		orgChecker := func(domain string) (string, error) {
+			return "test-org-id", nil
+		}
 		authenticator = Authenticator{
-			orgChecker:      orgLookup,
+			orgChecker:      orgChecker,
 			callbackHandler: func() (string, error) { return "authorizationCode", nil },
 			tokenRequester: func(authConfig astro.AuthConfig, verifier, code string) (Result, error) {
 				return Result{

cloud/deploy/deploy.go

@@ -209,8 +209,10 @@ func Deploy(deployInput InputDeploy, client astro.Client) error { //nolint
 		return nil
 	}
 
-	deploymentURL := "cloud." + domain + "/" + deployInfo.workspaceID + "/deployments/" + deployInfo.deploymentID + "/analytics"
-
+	deploymentURL, err := deployment.GetDeploymentURL(deployInfo.deploymentID, deployInfo.workspaceID)
+	if err != nil {
+		return err
+	}
 	if deployInput.Dags {
 		if len(dagFiles) == 0 && config.CFG.ShowWarnings.GetBool() {
 			i, _ := input.Confirm("Warning: No DAGs found. This will delete any existing DAGs. Are you sure you want to deploy?")

cloud/deployment/deployment.go

@@ -13,6 +13,7 @@ import (
 	astro "github.com/astronomer/astro-cli/astro-client"
 	"github.com/astronomer/astro-cli/config"
 	"github.com/astronomer/astro-cli/pkg/ansi"
+	"github.com/astronomer/astro-cli/pkg/domainutil"
 	"github.com/astronomer/astro-cli/pkg/httputil"
 	"github.com/astronomer/astro-cli/pkg/input"
 	"github.com/astronomer/astro-cli/pkg/printutil"
@@ -266,13 +267,10 @@ func createOutput(workspaceID string, d *astro.Deployment) error {
 
 	tab.AddRow([]string{d.Label, d.ReleaseName, workspaceID, d.Cluster.ID, d.ID, currentTag, runtimeVersionText, strconv.FormatBool(d.DagDeployEnabled)}, false)
 
-	c, err := config.GetCurrentContext()
+	deploymentURL, err := GetDeploymentURL(d.ID, workspaceID)
 	if err != nil {
 		return err
 	}
-
-	deploymentURL := "cloud." + c.Domain + "/" + workspaceID + "/deployments/" + d.ID + "/analytics"
-
 	tab.SuccessMsg = fmt.Sprintf("\n Successfully created Deployment: %s", ansi.Bold(d.Label)) +
 		"\n Deployment can be accessed at the following URLs \n" +
 		fmt.Sprintf("\n Deployment Dashboard: %s", ansi.Bold(deploymentURL)) +
@@ -714,6 +712,12 @@ func GetDeploymentURL(deploymentID, workspaceID string) (string, error) {
 	if err != nil {
 		return "", err
 	}
-	deploymentURL = "cloud." + ctx.Domain + "/" + workspaceID + "/deployments/" + deploymentID + "/analytics"
+	switch ctx.Domain {
+	case domainutil.LocalDomain:
+		deploymentURL = ctx.Domain + ":5000/" + workspaceID + "/deployments/" + deploymentID + "/analytics"
+	default:
+		_, domain := domainutil.GetPRSubDomain(ctx.Domain)
+		deploymentURL = "cloud." + domain + "/" + workspaceID + "/deployments/" + deploymentID + "/analytics"
+	}
 	return deploymentURL, nil
 }

cloud/deployment/deployment_test.go

@@ -970,9 +970,16 @@ func TestGetDeploymentURL(t *testing.T) {
 		assert.NoError(t, err)
 		assert.Equal(t, expectedURL, actualURL)
 	})
+	t.Run("returns deploymentURL for pr preview environment", func(t *testing.T) {
+		testUtil.InitTestConfig(testUtil.CloudPrPreview)
+		expectedURL := "cloud.astronomer-dev.io/workspace-id/deployments/deployment-id/analytics"
+		actualURL, err := GetDeploymentURL(deploymentID, workspaceID)
+		assert.NoError(t, err)
+		assert.Equal(t, expectedURL, actualURL)
+	})
 	t.Run("returns deploymentURL for local environment", func(t *testing.T) {
 		testUtil.InitTestConfig(testUtil.LocalPlatform)
-		expectedURL := "cloud.localhost/workspace-id/deployments/deployment-id/analytics"
+		expectedURL := "localhost:5000/workspace-id/deployments/deployment-id/analytics"
 		actualURL, err := GetDeploymentURL(deploymentID, workspaceID)
 		assert.NoError(t, err)
 		assert.Equal(t, expectedURL, actualURL)

cloud/deployment/inspect/inspect_test.go

@@ -323,7 +323,7 @@ func TestGetDeploymentInspectInfo(t *testing.T) {
 	t.Run("returns deployment Info for the requested local deployment", func(t *testing.T) {
 		var actualDeploymentInfo deploymentInfo
 		testUtil.InitTestConfig(testUtil.LocalPlatform)
-		expectedCloudDomainURL := "cloud.localhost/" + sourceDeployment.Workspace.ID +
+		expectedCloudDomainURL := "localhost:5000/" + sourceDeployment.Workspace.ID +
 			"/deployments/" + sourceDeployment.ID + "/analytics"
 		expectedDeploymentInfo := deploymentInfo{
 			DeploymentID:   sourceDeployment.ID,

cloud/organization/organization.go

@@ -13,6 +13,7 @@ import (
 	astro "github.com/astronomer/astro-cli/astro-client"
 	"github.com/astronomer/astro-cli/cloud/auth"
 	"github.com/astronomer/astro-cli/config"
+	"github.com/astronomer/astro-cli/pkg/domainutil"
 	"github.com/astronomer/astro-cli/pkg/httputil"
 	"github.com/astronomer/astro-cli/pkg/input"
 	"github.com/astronomer/astro-cli/pkg/printutil"
@@ -47,11 +48,10 @@ func newTableOut() *printutil.Table {
 // TODO use astro.go wrapper around REST client
 func listOrganizations(c *config.Context) ([]OrgRes, error) {
 	var orgDomain string
-	if c.Domain == "localhost" {
-		orgDomain = config.CFG.LocalCore.GetString() + "/organizations"
-	} else {
-		orgDomain = "https://api." + c.Domain + "/v1alpha1/organizations"
-	}
+	withOutCloud := domainutil.FormatDomain(c.Domain)
+	// we use core api for this
+	orgDomain = domainutil.GetURLToEndpoint("https", withOutCloud, "v1alpha1/organizations")
+	orgDomain = domainutil.TransformToCoreAPIEndpoint(orgDomain)
 	authToken := c.Token
 	ctx := context.Background()
 	doOptions := &httputil.DoOptions{

cmd/auth.go

@@ -3,6 +3,8 @@ package cmd
 import (
 	"io"
 
+	"github.com/astronomer/astro-cli/pkg/domainutil"
+
 	astro "github.com/astronomer/astro-cli/astro-client"
 	cloudAuth "github.com/astronomer/astro-cli/cloud/auth"
 	"github.com/astronomer/astro-cli/context"
@@ -66,7 +68,7 @@ func login(cmd *cobra.Command, args []string, astroClient astro.Client, out io.W
 	ctx, err := context.GetCurrentContext()
 	if err != nil || ctx.Domain == "" {
 		// Default case when no domain is passed, and error getting current context
-		return cloudLogin(cloudAuth.Domain, "", token, astroClient, out, shouldDisplayLoginLink)
+		return cloudLogin(domainutil.DefaultDomain, "", token, astroClient, out, shouldDisplayLoginLink)
 	} else if context.IsCloudDomain(ctx.Domain) {
 		return cloudLogin(ctx.Domain, "", token, astroClient, out, shouldDisplayLoginLink)
 	}