From 5fff911202ae24659a4b806f007aecd3e42b28f2 Mon Sep 17 00:00:00 2001
From: filip <filip.heden@debricked.com>
Date: Mon, 4 Nov 2024 15:41:16 +0100
Subject: [PATCH] Let viper handle env variable for token

---
 internal/client/deb_client.go      | 22 +++----------
 internal/client/deb_client_test.go | 50 +-----------------------------
 internal/client/request.go         |  7 +++--
 internal/cmd/root/root.go          |  9 ++++--
 internal/cmd/root/root_test.go     |  2 +-
 5 files changed, 16 insertions(+), 74 deletions(-)

diff --git a/internal/client/deb_client.go b/internal/client/deb_client.go
index d003eba3..a1714cef 100644
--- a/internal/client/deb_client.go
+++ b/internal/client/deb_client.go
@@ -41,13 +41,13 @@ func NewDebClient(accessToken *string, httpClient IClient) *DebClient {
 	if len(host) == 0 {
 		host = DefaultDebrickedUri
 	}
-	authenticator := auth.NewDebrickedAuthenticator(host)
+
 	return &DebClient{
 		host:          &host,
 		httpClient:    httpClient,
-		accessToken:   initAccessToken(accessToken),
+		accessToken:   accessToken,
 		jwtToken:      "",
-		authenticator: authenticator,
+		authenticator: auth.NewDebrickedAuthenticator(host),
 	}
 }
 
@@ -68,27 +68,13 @@ func (debClient *DebClient) Get(uri string, format string) (*http.Response, erro
 }
 
 func (debClient *DebClient) SetAccessToken(accessToken *string) {
-	debClient.accessToken = initAccessToken(accessToken)
+	debClient.accessToken = accessToken
 }
 
 func (debClient *DebClient) Authenticator() auth.IAuthenticator {
 	return debClient.authenticator
 }
 
-func initAccessToken(accessToken *string) *string {
-	if accessToken == nil {
-		accessToken = new(string)
-	}
-	if len(*accessToken) == 0 {
-		ok := false
-		*accessToken, ok = os.LookupEnv("DEBRICKED_TOKEN")
-		if !ok {
-			accessToken = nil
-		}
-	}
-	return accessToken
-}
-
 type BillingPlan struct {
 	SCA    string `json:"sca"`
 	Select string `json:"select"`
diff --git a/internal/client/deb_client_test.go b/internal/client/deb_client_test.go
index 99e350ca..5097ee95 100644
--- a/internal/client/deb_client_test.go
+++ b/internal/client/deb_client_test.go
@@ -36,37 +36,6 @@ func TestNewDebClientWithNilToken(t *testing.T) {
 	}
 }
 
-const debrickedTknEnvVar = "DEBRICKED_TOKEN"
-
-func TestNewDebClientWithTokenEnvVariable(t *testing.T) {
-	envVarTkn := "env-tkn"
-	oldEnvValue, ok := os.LookupEnv(debrickedTknEnvVar)
-	err := os.Setenv(debrickedTknEnvVar, "env-tkn")
-	if err != nil {
-		t.Fatalf("failed to set env var %s", debrickedTknEnvVar)
-	}
-	defer func(key, value string, ok bool) {
-		var err error = nil
-		if ok {
-			err = os.Setenv(key, value)
-		} else {
-			err = os.Unsetenv(key)
-		}
-		if err != nil {
-			t.Fatalf("failed to reset env var %s", debrickedTknEnvVar)
-		}
-	}(debrickedTknEnvVar, oldEnvValue, ok)
-
-	accessToken := ""
-	debClient := NewDebClient(&accessToken, nil)
-	if *debClient.host != DefaultDebrickedUri {
-		t.Error("failed to assert that host was set properly")
-	}
-	if *debClient.accessToken != envVarTkn {
-		t.Errorf("failed to assert that access token was set to %s. Got %s", envVarTkn, *debClient.accessToken)
-	}
-}
-
 func TestNewDebClientWithWithURI(t *testing.T) {
 	accessToken := ""
 	os.Setenv("DEBRICKED_URI", "https://subdomain.debricked.com")
@@ -215,24 +184,7 @@ func TestAuthenticateCachedToken(t *testing.T) {
 		authenticator: testdataAuth.MockAuthenticator{},
 	}
 
-	oldEnvValue, ok := os.LookupEnv(debrickedTknEnvVar)
-	err := os.Unsetenv(debrickedTknEnvVar)
-	if err != nil {
-		t.Fatalf("failed to set env var %s", debrickedTknEnvVar)
-	}
-	defer func(key, value string, ok bool) {
-		var err error = nil
-		if ok {
-			err = os.Setenv(key, value)
-		} else {
-			err = os.Unsetenv(key)
-		}
-		if err != nil {
-			t.Fatalf("failed to reset env var %s", debrickedTknEnvVar)
-		}
-	}(debrickedTknEnvVar, oldEnvValue, ok)
-
-	err = client.authenticate()
+	err := client.authenticate()
 	if err != nil {
 		t.Fatal("failed to assert that no error occurred")
 	}
diff --git a/internal/client/request.go b/internal/client/request.go
index f87d0309..d352377f 100644
--- a/internal/client/request.go
+++ b/internal/client/request.go
@@ -117,9 +117,10 @@ type errorMessage struct {
 }
 
 func (debClient *DebClient) authenticate() error {
-
-	if debClient.accessToken != nil {
-		return debClient.authenticateExplicitToken()
+	if debClient.accessToken != nil { // To avoid segfault
+		if len(*debClient.accessToken) != 0 {
+			return debClient.authenticateExplicitToken()
+		}
 	}
 
 	return debClient.authenticateCachedToken()
diff --git a/internal/cmd/root/root.go b/internal/cmd/root/root.go
index 51a9c481..0ace847b 100644
--- a/internal/cmd/root/root.go
+++ b/internal/cmd/root/root.go
@@ -15,7 +15,8 @@ import (
 
 var accessToken string
 
-const AccessTokenFlag = "access-token"
+const AccessTokenFlag = "token"
+const OldAccessTokenFlag = "access-token"
 
 func NewRootCmd(version string, container *wire.CliContainer) *cobra.Command {
 	rootCmd := &cobra.Command{
@@ -29,15 +30,17 @@ Complete documentation is available at https://docs.debricked.com/tools-and-inte
 		Version: version,
 	}
 	viper.SetEnvPrefix("DEBRICKED")
-	viper.MustBindEnv(AccessTokenFlag)
+	viper.AutomaticEnv()
 	rootCmd.PersistentFlags().StringVarP(
 		&accessToken,
 		AccessTokenFlag,
 		"t",
-		"",
+		viper.GetString(AccessTokenFlag),
 		`Debricked access token. 
 Read more: https://docs.debricked.com/product/administration/generate-access-token`,
 	)
+	viper.MustBindEnv(AccessTokenFlag)
+	viper.RegisterAlias(OldAccessTokenFlag, AccessTokenFlag) // To avoid breaking change
 
 	var debClient = container.DebClient()
 	debClient.SetAccessToken(&accessToken)
diff --git a/internal/cmd/root/root_test.go b/internal/cmd/root/root_test.go
index 52e32cb8..0fd0cff4 100644
--- a/internal/cmd/root/root_test.go
+++ b/internal/cmd/root/root_test.go
@@ -35,7 +35,7 @@ func TestNewRootCmd(t *testing.T) {
 		}
 	}
 	assert.Truef(t, match, "failed to assert that flag was present: "+AccessTokenFlag)
-	assert.Len(t, viperKeys, 22)
+	assert.Len(t, viperKeys, 23)
 }
 
 func TestPreRun(t *testing.T) {