feat(az): display account info

docs/docs/segment-az.mdx

@@ -33,6 +33,7 @@ To enable this, set `$env:AZ_ENABLED = $true` in your `$PROFILE`.
 
 ## Properties
 
-- info_separator: `string` - text/icon to put in between the subscription name and ID - defaults to ` | `
-- display_id: `boolean` - display the subscription ID or not - defaults to `false`
+- display_account: `boolean` - display the subscription account name or not - defaults to `false`
 - display_name: `boolean` - display the subscription name or not - defaults to `true`
+- display_id: `boolean` - display the subscription ID or not - defaults to `false`
+- info_separator: `string` - text/icon to put in between the values - defaults to ` | `

src/init/omp.ps1

@@ -47,10 +47,11 @@ function global:Initialize-ModuleSupport {
 
     if ($env:AZ_ENABLED -eq $true) {
         try {
-            $subscription = Get-AzContext | Select-Object -ExpandProperty "Subscription" | Select-Object "Name", "Id"
+            $subscription = Get-AzContext | Select-Object -ExpandProperty "Subscription" | Select-Object "Name", "Id", "Account"
             if ($null -ne $subscription) {
                 $env:AZ_SUBSCRIPTION_NAME = $subscription.Name
                 $env:AZ_SUBSCRIPTION_ID = $subscription.Id
+                $env:AZ_SUBSCRIPTION_ACCOUNT = $subscription.Account
             }
         }
         catch {}

src/segment_az.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"fmt"
 	"strings"
 )
 
@@ -10,6 +9,9 @@ type az struct {
 	env       environmentInfo
 	name      string
 	id        string
+	account   string
+	builder   strings.Builder
+	separator string
 }
 
 const (
@@ -19,6 +21,8 @@ const (
 	DisplaySubscriptionID Property = "display_id"
 	// DisplaySubscriptionName hides or shows the subscription display name
 	DisplaySubscriptionName Property = "display_name"
+	// DisplaySubscriptionAccount hides or shows the subscription account name
+	DisplaySubscriptionAccount Property = "display_account"
 
 	updateConsentNeeded = "Do you want to continue?"
 	updateMessage       = "AZ CLI: Update needed!"
@@ -27,12 +31,27 @@ const (
 )
 
 func (a *az) string() string {
-	separator := ""
-	if a.idEnabled() && a.nameEnabled() {
-		separator = a.props.getString(SubscriptionInfoSeparator, " | ")
+	a.separator = a.props.getString(SubscriptionInfoSeparator, " | ")
+	writeValue := func(value string) {
+		if len(value) == 0 {
+			return
+		}
+		if a.builder.Len() > 0 {
+			a.builder.WriteString(a.separator)
+		}
+		a.builder.WriteString(value)
+	}
+	if a.props.getBool(DisplaySubscriptionAccount, false) {
+		writeValue(a.account)
+	}
+	if a.props.getBool(DisplaySubscriptionName, true) {
+		writeValue(a.name)
+	}
+	if a.props.getBool(DisplaySubscriptionID, false) {
+		writeValue(a.id)
 	}
 
-	return fmt.Sprintf("%s%s%s", a.getName(), separator, a.getID())
+	return a.builder.String()
 }
 
 func (a *az) init(props *properties, env environmentInfo) {
@@ -41,75 +60,50 @@ func (a *az) init(props *properties, env environmentInfo) {
 }
 
 func (a *az) enabled() bool {
-	var enabled bool
-	a.name, a.id, enabled = a.getFromEnvVars()
-	if enabled {
-		return enabled
+	if a.getFromEnvVars() {
+		return true
 	}
 
-	a.name, a.id, enabled = a.getFromAzCli()
-	return enabled
+	return a.getFromAzCli()
 }
 
-func (a *az) getFromEnvVars() (string, string, bool) {
-	name := a.env.getenv("AZ_SUBSCRIPTION_NAME")
-	id := a.env.getenv("AZ_SUBSCRIPTION_ID")
+func (a *az) getFromEnvVars() bool {
+	a.name = a.env.getenv("AZ_SUBSCRIPTION_NAME")
+	a.id = a.env.getenv("AZ_SUBSCRIPTION_ID")
+	a.account = a.env.getenv("AZ_SUBSCRIPTION_ID")
 
-	if name == "" && id == "" {
-		return "", "", false
+	if a.name == "" && a.id == "" {
+		return false
 	}
 
-	return name, id, true
+	return true
 }
 
-func (a *az) getFromAzCli() (string, string, bool) {
+func (a *az) getFromAzCli() bool {
 	cmd := "az"
-	if (!a.idEnabled() && !a.nameEnabled()) || !a.env.hasCommand(cmd) {
-		return "", "", false
+	if !a.env.hasCommand(cmd) {
+		return false
 	}
 
-	output, _ := a.env.runCommand(cmd, "account", "show", "--query=[name,id]", "-o=tsv")
-	if output == "" {
-		return "", "", false
+	output, _ := a.env.runCommand(cmd, "account", "show", "--query=[name,id,user.name]", "-o=tsv")
+	if len(output) == 0 {
+		return false
 	}
 
 	if strings.Contains(output, updateConsentNeeded) {
 		a.props.foreground = updateForeground
 		a.props.background = updateBackground
-		return updateMessage, "", true
+		a.name = updateMessage
+		return true
 	}
 
 	splittedOutput := strings.Split(output, "\n")
-	if len(splittedOutput) < 2 {
-		return "", "", false
+	if len(splittedOutput) < 3 {
+		return false
 	}
 
-	name := strings.TrimSpace(splittedOutput[0])
-	id := strings.TrimSpace(splittedOutput[1])
-
-	return name, id, true
-}
-
-func (a *az) getID() string {
-	if !a.idEnabled() {
-		return ""
-	}
-
-	return a.id
-}
-
-func (a *az) getName() string {
-	if !a.nameEnabled() {
-		return ""
-	}
-
-	return a.name
-}
-
-func (a *az) idEnabled() bool {
-	return a.props.getBool(DisplaySubscriptionID, false)
-}
-
-func (a *az) nameEnabled() bool {
-	return a.props.getBool(DisplaySubscriptionName, true)
+	a.name = strings.TrimSpace(splittedOutput[0])
+	a.id = strings.TrimSpace(splittedOutput[1])
+	a.account = strings.TrimSpace(splittedOutput[2])
+	return true
 }

src/segment_az_test.go

@@ -14,13 +14,29 @@ func TestAzSegment(t *testing.T) {
 		ExpectedString  string
 		EnvSubName      string
 		EnvSubID        string
+		EnvSubAccount   string
 		CliExists       bool
 		CliSubName      string
 		CliSubID        string
+		CliSubAccount   string
 		InfoSeparator   string
 		DisplayID       bool
 		DisplayName     bool
+		DisplayAccount  bool
 	}{
+		{
+			Case:            "print only account",
+			ExpectedEnabled: true,
+			ExpectedString:  "foobar",
+			CliExists:       true,
+			CliSubName:      "foo",
+			CliSubID:        "bar",
+			CliSubAccount:   "foobar",
+			InfoSeparator:   "$",
+			DisplayID:       false,
+			DisplayName:     false,
+			DisplayAccount:  true,
+		},
 		{
 			Case:            "envvars present",
 			ExpectedEnabled: true,
@@ -35,7 +51,7 @@ func TestAzSegment(t *testing.T) {
 		{
 			Case:            "envvar name present",
 			ExpectedEnabled: true,
-			ExpectedString:  "foo$",
+			ExpectedString:  "foo",
 			EnvSubName:      "foo",
 			CliExists:       false,
 			InfoSeparator:   "$",
@@ -45,7 +61,7 @@ func TestAzSegment(t *testing.T) {
 		{
 			Case:            "envvar id present",
 			ExpectedEnabled: true,
-			ExpectedString:  "$bar",
+			ExpectedString:  "bar",
 			EnvSubID:        "bar",
 			CliExists:       false,
 			InfoSeparator:   "$",
@@ -55,7 +71,7 @@ func TestAzSegment(t *testing.T) {
 		{
 			Case:            "cli not found",
 			ExpectedEnabled: false,
-			ExpectedString:  "$",
+			ExpectedString:  "",
 			CliExists:       false,
 			InfoSeparator:   "$",
 			DisplayID:       true,
@@ -96,13 +112,11 @@ func TestAzSegment(t *testing.T) {
 		},
 		{
 			Case:            "print none",
-			ExpectedEnabled: false,
+			ExpectedEnabled: true,
 			CliExists:       true,
 			CliSubName:      "foo",
 			CliSubID:        "bar",
 			InfoSeparator:   "$",
-			DisplayID:       false,
-			DisplayName:     false,
 		},
 		{
 			Case:            "update needed",
@@ -113,6 +127,16 @@ func TestAzSegment(t *testing.T) {
 			DisplayID:       false,
 			DisplayName:     true,
 		},
+		{
+			Case:            "account info",
+			ExpectedEnabled: true,
+			ExpectedString:  updateMessage,
+			CliExists:       true,
+			CliSubName:      "Do you want to continue? (Y/n): Visual Studio Enterprise",
+			DisplayID:       false,
+			DisplayName:     true,
+			DisplayAccount:  true,
+		},
 	}
 
 	for _, tc := range cases {
@@ -120,12 +144,16 @@ func TestAzSegment(t *testing.T) {
 		env.On("getenv", "AZ_SUBSCRIPTION_NAME").Return(tc.EnvSubName)
 		env.On("getenv", "AZ_SUBSCRIPTION_ID").Return(tc.EnvSubID)
 		env.On("hasCommand", "az").Return(tc.CliExists)
-		env.On("runCommand", "az", []string{"account", "show", "--query=[name,id]", "-o=tsv"}).Return(fmt.Sprintf("%s\n%s\n", tc.CliSubName, tc.CliSubID), nil)
+		env.On("runCommand", "az", []string{"account", "show", "--query=[name,id,user.name]", "-o=tsv"}).Return(
+			fmt.Sprintf("%s\n%s\n%s\n", tc.CliSubName, tc.CliSubID, tc.CliSubAccount),
+			nil,
+		)
 		props := &properties{
 			values: map[Property]interface{}{
 				SubscriptionInfoSeparator:  tc.InfoSeparator,
 				DisplaySubscriptionID:      tc.DisplayID,
 				DisplaySubscriptionName:    tc.DisplayName,
+				DisplaySubscriptionAccount: tc.DisplayAccount,
 			},
 		}