diff --git a/docs/docs/segment-az.mdx b/docs/docs/segment-az.mdx index ab802bb1..58cbd7d1 100644 --- a/docs/docs/segment-az.mdx +++ b/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 ` | ` diff --git a/src/init/omp.ps1 b/src/init/omp.ps1 index cb73d595..5592ce5b 100644 --- a/src/init/omp.ps1 +++ b/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 {} diff --git a/src/segment_az.go b/src/segment_az.go index 4b2afb5e..088f18a0 100644 --- a/src/segment_az.go +++ b/src/segment_az.go @@ -1,15 +1,17 @@ package main import ( - "fmt" "strings" ) type az struct { - props *properties - env environmentInfo - name string - id string + props *properties + 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 } diff --git a/src/segment_az_test.go b/src/segment_az_test.go index c9656744..6f6087e5 100644 --- a/src/segment_az_test.go +++ b/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, + SubscriptionInfoSeparator: tc.InfoSeparator, + DisplaySubscriptionID: tc.DisplayID, + DisplaySubscriptionName: tc.DisplayName, + DisplaySubscriptionAccount: tc.DisplayAccount, }, }