From b1466cf90cba758c6cfb1e8b27a6fc5217148912 Mon Sep 17 00:00:00 2001
From: "Yun Lu (MSFT)" <53427149+Luyunmt@users.noreply.github.com>
Date: Tue, 7 Apr 2020 03:58:34 +0800
Subject: [PATCH] Add EnvVar AZURE_AUTHORITY_HOST (#8074)

* add env authorityhost

* add AuthorityHost tests

* Unset env variable after AuthorityTest to avoid other tests fail

* add default AuthorityHost test

Co-authored-by: wantedfast <inaccuracy@gmail.com>
Co-authored-by: Geralt Wang <v-xinlwa@microsoft.com>
---
 sdk/azidentity/azidentity.go      | 10 ++++--
 sdk/azidentity/azidentity_test.go | 57 +++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/sdk/azidentity/azidentity.go b/sdk/azidentity/azidentity.go
index 08df5821d86b..da95a058d24b 100644
--- a/sdk/azidentity/azidentity.go
+++ b/sdk/azidentity/azidentity.go
@@ -6,6 +6,7 @@ package azidentity
 import (
 	"net/http"
 	"net/url"
+	"os"
 	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
@@ -117,8 +118,13 @@ type TokenCredentialOptions struct {
 // NewIdentityClientOptions initializes an instance of IdentityClientOptions with default settings
 // NewIdentityClientOptions initializes an instance of IdentityClientOptions with default settings
 func (c *TokenCredentialOptions) setDefaultValues() (*TokenCredentialOptions, error) {
+	authorityHost := defaultAuthorityHost
+	if envAuthorityHost := os.Getenv("AZURE_AUTHORITY_HOST"); envAuthorityHost != "" {
+		authorityHost = envAuthorityHost
+	}
+
 	if c == nil {
-		defaultAuthorityHostURL, err := url.Parse(defaultAuthorityHost)
+		defaultAuthorityHostURL, err := url.Parse(authorityHost)
 		if err != nil {
 			return nil, err
 		}
@@ -126,7 +132,7 @@ func (c *TokenCredentialOptions) setDefaultValues() (*TokenCredentialOptions, er
 	}
 
 	if c.AuthorityHost == nil {
-		defaultAuthorityHostURL, err := url.Parse(defaultAuthorityHost)
+		defaultAuthorityHostURL, err := url.Parse(authorityHost)
 		if err != nil {
 			return nil, err
 		}
diff --git a/sdk/azidentity/azidentity_test.go b/sdk/azidentity/azidentity_test.go
index 5fb5761819d4..ed6b0135132e 100644
--- a/sdk/azidentity/azidentity_test.go
+++ b/sdk/azidentity/azidentity_test.go
@@ -5,11 +5,17 @@ package azidentity
 
 import (
 	"net/url"
+	"os"
 	"testing"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 )
 
+const (
+	envHostString    = "https://mock.com/"
+	customHostString = "https://custommock.com/"
+)
+
 func Test_AuthorityHost_Parse(t *testing.T) {
 	_, err := url.Parse(defaultAuthorityHost)
 	if err != nil {
@@ -27,3 +33,54 @@ func Test_NonNilTokenCredentialOptsNilAuthorityHost(t *testing.T) {
 		t.Fatalf("Did not set default authority host")
 	}
 }
+
+func Test_SetEnvAuthorityHost(t *testing.T) {
+	err := os.Setenv("AZURE_AUTHORITY_HOST", envHostString)
+	if err != nil {
+		t.Fatalf("Unexpected error when initializing environment variables: %v", err)
+	}
+
+	opts := &TokenCredentialOptions{}
+	opts, err = opts.setDefaultValues()
+	if opts.AuthorityHost.String() != envHostString {
+		t.Fatalf("Unexpected error when get host from environment vairable: %v", err)
+	}
+
+	// Unset that host environment vairable to avoid other tests failed.
+	err = os.Unsetenv("AZURE_AUTHORITY_HOST")
+	if err != nil {
+		t.Fatalf("Unexpected error when unset environment vairable: %v", err)
+	}
+}
+
+func Test_CustomAuthorityHost(t *testing.T) {
+	err := os.Setenv("AZURE_AUTHORITY_HOST", envHostString)
+	if err != nil {
+		t.Fatalf("Unexpected error when initializing environment variables: %v", err)
+	}
+
+	customHost, err := url.Parse(customHostString)
+	if err != nil {
+		t.Fatalf("Received an error: %v", err)
+	}
+
+	opts := &TokenCredentialOptions{AuthorityHost: customHost}
+	opts, err = opts.setDefaultValues()
+	if opts.AuthorityHost.String() != customHostString {
+		t.Fatalf("Unexpected error when get host from environment vairable: %v", err)
+	}
+
+	// Unset that host environment vairable to avoid other tests failed.
+	err = os.Unsetenv("AZURE_AUTHORITY_HOST")
+	if err != nil {
+		t.Fatalf("Unexpected error when unset environment vairable: %v", err)
+	}
+}
+
+func Test_DefaultAuthorityHost(t *testing.T) {
+	opts := &TokenCredentialOptions{}
+	opts, err := opts.setDefaultValues()
+	if opts.AuthorityHost.String() != defaultAuthorityHost {
+		t.Fatalf("Unexpected error when set default AuthorityHost: %v", err)
+	}
+}