Allow localhost on u2m auth

credentials/u2m/account_oauth_argument.go

@@ -2,7 +2,6 @@ package u2m
 
 import (
 	"fmt"
-	"strings"
 )
 
 // AccountOAuthArgument is an interface that provides the necessary information
@@ -28,11 +27,8 @@ var _ AccountOAuthArgument = BasicAccountOAuthArgument{}
 
 // NewBasicAccountOAuthArgument creates a new BasicAccountOAuthArgument.
 func NewBasicAccountOAuthArgument(accountsHost, accountID string) (BasicAccountOAuthArgument, error) {
-	if !strings.HasPrefix(accountsHost, "https://") {
-		return BasicAccountOAuthArgument{}, fmt.Errorf("accountsHost must start with 'https://': %s", accountsHost)
-	}
-	if strings.HasSuffix(accountsHost, "/") {
-		return BasicAccountOAuthArgument{}, fmt.Errorf("accountsHost must not have a trailing slash: %s", accountsHost)
+	if err := validateHost(accountsHost); err != nil {
+		return BasicAccountOAuthArgument{}, err
 	}
 	return BasicAccountOAuthArgument{accountHost: accountsHost, accountID: accountID}, nil
 }

credentials/u2m/workspace_oauth_argument.go

@@ -22,13 +22,25 @@ type BasicWorkspaceOAuthArgument struct {
 	host string
 }
 
-// NewBasicWorkspaceOAuthArgument creates a new BasicWorkspaceOAuthArgument.
-func NewBasicWorkspaceOAuthArgument(host string) (BasicWorkspaceOAuthArgument, error) {
+func validateHost(host string) error {
+	// Allow http for localhost. This is necessary for local end to end testing
+	// of the `databricks auth login` command using a test server on localhost.
+	if strings.HasPrefix(host, "http://127.0.0.1") {
+		return nil
+	}
 	if !strings.HasPrefix(host, "https://") {
-		return BasicWorkspaceOAuthArgument{}, fmt.Errorf("host must start with 'https://': %s", host)
+		return fmt.Errorf("host must start with 'https://': %s", host)
 	}
 	if strings.HasSuffix(host, "/") {
-		return BasicWorkspaceOAuthArgument{}, fmt.Errorf("host must not have a trailing slash: %s", host)
+		return fmt.Errorf("host must not have a trailing slash: %s", host)
+	}
+	return nil
+}
+
+// NewBasicWorkspaceOAuthArgument creates a new BasicWorkspaceOAuthArgument.
+func NewBasicWorkspaceOAuthArgument(host string) (BasicWorkspaceOAuthArgument, error) {
+	if err := validateHost(host); err != nil {
+		return BasicWorkspaceOAuthArgument{}, err
 	}
 	return BasicWorkspaceOAuthArgument{host: host}, nil
 }

credentials/u2m/workspace_oauth_argument_test.go

@@ -0,0 +1,46 @@
+package u2m
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestValidateHost(t *testing.T) {
+	tests := []struct {
+		host string
+		want string
+	}{
+		// Valid hosts
+		{"https://some-host.com", ""},
+		{"http://127.0.0.1", ""},
+		{"http://127.0.0.1:5656", ""},
+
+		// Invalid hosts
+		{"http://some-host.com", "host must start with 'https://': http://some-host.com"},
+		{"https://some-host.com/", "host must not have a trailing slash: https://some-host.com/"},
+	}
+
+	for _, test := range tests {
+		err := validateHost(test.host)
+		if test.want == "" {
+			assert.NoError(t, err)
+		} else {
+			assert.EqualError(t, err, test.want)
+		}
+
+		_, err = NewBasicWorkspaceOAuthArgument(test.host)
+		if test.want == "" {
+			assert.NoError(t, err)
+		} else {
+			assert.EqualError(t, err, test.want)
+		}
+
+		_, err = NewBasicAccountOAuthArgument(test.host, "123")
+		if test.want == "" {
+			assert.NoError(t, err)
+		} else {
+			assert.EqualError(t, err, test.want)
+		}
+	}
+}