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
astronomer · Nov 11, 2022 · d44b351 · d44b351
1 parent fba8ebc
commit d44b351
Show file tree

Hide file tree

Showing 9 changed files with 377 additions and 154 deletions.
diff --git a/cloud/auth/auth.go b/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 {
@@ -358,7 +346,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
@@ -382,7 +370,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
@@ -433,46 +420,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") {
-		domain = strings.Replace(domain, "cloud.", "", 1)
-	} else if domain == "" {
-		domain = Domain
-	}
-
-	return domain
-}
-
 func ValidateDomain(domain string) (astro.AuthConfig, error) {
 	var (
 		authConfig  astro.AuthConfig
-		prNum, addr string
+		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'")
 	}
 
-	// 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",
-			domain,
-		)
-	}
+	addr = domainutil.GetURLToEndpoint("https", domain, authConfigEndpoint)
+
 	ctx := http_context.Background()
 	doOptions := &httputil.DoOptions{
 		Context: ctx,

diff --git a/cloud/auth/auth_test.go b/cloud/auth/auth_test.go
@@ -89,13 +89,27 @@ func Test_validateDomain(t *testing.T) {
 		"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"
+		// 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, "GWCArPCS6pgFZbhnwx3WCDaJLBcvdFaV")
-		assert.Equal(t, actual.Audience, "astronomer-ee")
-		assert.Equal(t, actual.DomainURL, "https://astronomer-sandbox.us.auth0.com")
+		assert.Equal(t, actual.ClientID, mockResponse.ClientID)
+		assert.Equal(t, actual.Audience, mockResponse.Audience)
+		assert.Equal(t, actual.DomainURL, mockResponse.DomainURL)
 	})
 }
 
@@ -467,6 +481,21 @@ func TestLogin(t *testing.T) {
 	})
 	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
@@ -486,7 +515,7 @@ func TestLogin(t *testing.T) {
 		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)
+		err = Login("pr5723.cloud.astronomer-dev.io", "", "", mockClient, os.Stdout, false)
 		assert.NoError(t, err)
 	})
 
@@ -551,8 +580,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{
@@ -580,8 +612,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{
@@ -720,30 +755,3 @@ 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)
-	})
-}
diff --git a/cloud/organization/organization.go b/cloud/organization/organization.go
@@ -12,6 +12,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"
@@ -45,11 +46,10 @@ func newTableOut() *printutil.Table {
 
 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{

diff --git a/cmd/auth.go b/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)
 	}

diff --git a/config/cloud.go b/config/cloud.go
@@ -5,9 +5,14 @@ import (
 	"io"
 	"strings"
 
+	"github.com/astronomer/astro-cli/pkg/domainutil"
 	"github.com/astronomer/astro-cli/pkg/printutil"
 )
 
+const (
+	graphqlEndpoint = "graphql"
+)
+
 // PrintCloudContext prints current context to stdOut
 func (c *Context) PrintCloudContext(out io.Writer) error {
 	context, err := c.GetContext()
@@ -75,15 +80,6 @@ func (c *Context) GetPublicAPIURL() string {
 		return CFG.LocalPublicAstro.GetString()
 	}
 
-	domain := c.Domain
-	if strings.Contains(domain, cloudDomain) {
-		splitDomain := strings.SplitN(domain, ".", splitNum) // This splits out 'cloud' from the domain string
-		domain = splitDomain[1]
-	}
-
-	return fmt.Sprintf(
-		"%s://api.%s/hub/graphql",
-		CFG.CloudAPIProtocol.GetString(),
-		domain,
-	)
+	domain := domainutil.FormatDomain(c.Domain)
+	return domainutil.GetURLToEndpoint(CFG.CloudAPIProtocol.GetString(), domain, graphqlEndpoint)
 }
diff --git a/context/context.go b/context/context.go
@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"github.com/astronomer/astro-cli/config"
+	"github.com/astronomer/astro-cli/pkg/domainutil"
 	"github.com/astronomer/astro-cli/pkg/input"
 	"github.com/astronomer/astro-cli/pkg/printutil"
 	"github.com/spf13/cobra"
@@ -16,8 +17,7 @@ var (
 	// CloudDomainRegex is used to differentiate cloud domain from software domain
 	// See https://github.com/astronomer/astrohub-cli/issues/7 for regexp rationale
 	// This will need to be handled as part of the permanent solution to issue #432
-	CloudDomainRegex = regexp.MustCompile(`(cloud\.|^)astronomer(?:(-dev|-stage|-perf))?\.io`)
-
+	CloudDomainRegex     = regexp.MustCompile(`(?:(pr\d{4,6})\.|^)(?:cloud\.|)astronomer(?:-(dev|stage|perf))?\.io$`)
 	contextDeleteWarnMsg = "Are you sure you want to delete currently used context: %s"
 	cancelCtxDeleteMsg   = "Canceling context delete..."
 	failCtxDeleteMsg     = "Error deleting context %s: "
@@ -30,12 +30,6 @@ var tab = printutil.Table{
 	ColorRowCode: [2]string{"\033[1;32m", "\033[0m"},
 }
 
-const (
-	DevCloudDomain   = "astronomer-dev.io"
-	StageCloudDomain = "astronomer-stage.io"
-	PerfCloudDomain  = "astronomer-perf.io"
-)
-
 // ContextExists checks to see if context exist in config
 func Exists(domain string) bool {
 	c := config.Context{Domain: domain}
@@ -166,10 +160,10 @@ func IsCloudContext() bool {
 
 // IsCloudDomain returns whether the given domain is related to cloud platform or not
 func IsCloudDomain(domain string) bool {
-	if domain == DevCloudDomain ||
-		domain == StageCloudDomain ||
-		domain == PerfCloudDomain ||
-		CloudDomainRegex.MatchString(domain) {
+	if CloudDomainRegex.MatchString(domain) {
+		return true
+	}
+	if domainutil.PRPreviewDomainRegex.MatchString(domain) {
 		return true
 	}