feat(az): support az pwsh module

docs/docs/segment-environment.md

@@ -29,7 +29,7 @@ make sure to do this after importing `go-my-posh` and you're good to go.
 function Set-EnvVar {
     $env:POSH=$(Get-Date)
 }
-New-Alias -Name 'Set-PoshContext' -Value 'Set-EnvVar' -Scope Global
+New-Alias -Name 'Set-PoshContext' -Value 'Set-EnvVar' -Scope Global -Force
 ```
 
 The segment will show when the value of the environment variable isn't empty.

src/init/omp.ps1

@@ -16,12 +16,23 @@ if (Test-Path $config) {
 
 function global:Set-PoshContext {}
 
-function global:Set-PoshGitStatus {
+function global:Initialize-ModuleSupport {
     if (Get-Module -Name "posh-git") {
         [Diagnostics.CodeAnalysis.SuppressMessageAttribute('PSProvideCommentHelp', '', Justification = 'Variable used later(not in this scope)')]
         $global:GitStatus = Get-GitStatus
         $env:POSH_GIT_STATUS = Write-GitStatus -Status $global:GitStatus
     }
+
+    if (Get-Module -ListAvailable -Name "Az.Accounts") {
+        try {
+            $subscription = Get-AzContext | Select-Object -ExpandProperty "Subscription" | Select-Object "Name", "Id"
+            if ($null -ne $subscription) {
+                $env:AZ_SUBSCRIPTION_NAME = $subscription.Name
+                $env:AZ_SUBSCRIPTION_ID = $subscription.Id
+            }
+        }
+        catch {}
+    }
 }
 
 [ScriptBlock]$Prompt = {
@@ -44,10 +55,10 @@ function global:Set-PoshGitStatus {
     $executionTime = -1
     $history = Get-History -ErrorAction Ignore -Count 1
     if ($null -ne $history -and $null -ne $history.EndExecutionTime -and $null -ne $history.StartExecutionTime -and $global:omp_lastHistoryId -ne $history.Id) {
-            $executionTime = ($history.EndExecutionTime - $history.StartExecutionTime).TotalMilliseconds
+        $executionTime = ($history.EndExecutionTime - $history.StartExecutionTime).TotalMilliseconds
-            $global:omp_lastHistoryId = $history.Id
+        $global:omp_lastHistoryId = $history.Id
     }
-    Set-PoshGitStatus
+    Initialize-ModuleSupport
     $omp = "::OMP::"
     $config = $global:PoshSettings.Theme
     $cleanPWD = $PWD.ProviderPath.TrimEnd("\")
@@ -77,7 +88,7 @@ function global:Export-PoshTheme {
         [string]
         $FilePath,
         [Parameter(Mandatory = $false)]
-        [ValidateSet('json','yaml','toml')]
+        [ValidateSet('json', 'yaml', 'toml')]
         [string]
         $Format = 'json'
     )

src/segment_az.go

@@ -36,24 +36,47 @@ 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
+	}
+
+	a.name, a.id, enabled = a.getFromAzCli()
+	return enabled
+}
+
+func (a *az) getFromEnvVars() (string, string, bool) {
+	name := a.env.getenv("AZ_SUBSCRIPTION_NAME")
+	id := a.env.getenv("AZ_SUBSCRIPTION_ID")
+
+	if name == "" && id == "" {
+		return "", "", false
+	}
+
+	return name, id, true
+}
+
+func (a *az) getFromAzCli() (string, string, bool) {
 	cmd := "az"
 	if (!a.idEnabled() && !a.nameEnabled()) || !a.env.hasCommand(cmd) {
-		return false
+		return "", "", false
 	}
 
 	output, _ := a.env.runCommand(cmd, "account", "show", "--query=[name,id]", "-o=tsv")
 	if output == "" {
-		return false
+		return "", "", false
 	}
 
 	splittedOutput := strings.Split(output, "\n")
 	if len(splittedOutput) < 2 {
-		return false
+		return "", "", false
 	}
 
-	a.name = strings.TrimSpace(splittedOutput[0])
+	name := strings.TrimSpace(splittedOutput[0])
-	a.id = strings.TrimSpace(splittedOutput[1])
+	id := strings.TrimSpace(splittedOutput[1])
-	return true
+
+	return name, id, true
 }
 
 func (a *az) getID() string {

src/segment_az_test.go

@@ -7,91 +7,121 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-type azArgs struct {
+func TestAzSegment(t *testing.T) {
-	enabled          bool
+	cases := []struct {
-	subscriptionName string
+		Case            string
-	subscriptionID   string
+		ExpectedEnabled bool
-	infoSeparator    string
+		ExpectedString  string
-	displayID        bool
+		EnvSubName      string
-	displayName      bool
+		EnvSubID        string
-}
+		CliExists       bool
-
+		CliSubName      string
-func bootStrapAzTest(args *azArgs) *az {
+		CliSubID        string
-	env := new(MockedEnvironment)
+		InfoSeparator   string
-	env.On("hasCommand", "az").Return(args.enabled)
+		DisplayID       bool
-	env.On("runCommand", "az", []string{"account", "show", "--query=[name,id]", "-o=tsv"}).Return(fmt.Sprintf("%s\n%s\n", args.subscriptionName, args.subscriptionID), nil)
+		DisplayName     bool
-	props := &properties{
+	}{
-		values: map[Property]interface{}{
+		{Case: "envvars present",
-			SubscriptionInfoSeparator: args.infoSeparator,
+			ExpectedEnabled: true,
-			DisplaySubscriptionID:     args.displayID,
+			ExpectedString:  "foo$bar",
-			DisplaySubscriptionName:   args.displayName,
+			EnvSubName:      "foo",
-		},
+			EnvSubID:        "bar",
+			CliExists:       false,
+			InfoSeparator:   "$",
+			DisplayID:       true,
+			DisplayName:     true},
+		{Case: "envvar name present",
+			ExpectedEnabled: true,
+			ExpectedString:  "foo$",
+			EnvSubName:      "foo",
+			EnvSubID:        "",
+			CliExists:       false,
+			InfoSeparator:   "$",
+			DisplayID:       true,
+			DisplayName:     true},
+		{Case: "envvar id present",
+			ExpectedEnabled: true,
+			ExpectedString:  "$bar",
+			EnvSubName:      "",
+			EnvSubID:        "bar",
+			CliExists:       false,
+			InfoSeparator:   "$",
+			DisplayID:       true,
+			DisplayName:     true},
+		{Case: "cli not found",
+			ExpectedEnabled: false,
+			ExpectedString:  "$",
+			EnvSubName:      "",
+			EnvSubID:        "",
+			CliExists:       false,
+			InfoSeparator:   "$",
+			DisplayID:       true,
+			DisplayName:     true},
+		{Case: "cli contains data",
+			ExpectedEnabled: true,
+			ExpectedString:  "foo$bar",
+			EnvSubName:      "",
+			EnvSubID:        "",
+			CliExists:       true,
+			CliSubName:      "foo",
+			CliSubID:        "bar",
+			InfoSeparator:   "$",
+			DisplayID:       true,
+			DisplayName:     true},
+		{Case: "print only name",
+			ExpectedEnabled: true,
+			ExpectedString:  "foo",
+			EnvSubName:      "",
+			EnvSubID:        "",
+			CliExists:       true,
+			CliSubName:      "foo",
+			CliSubID:        "bar",
+			InfoSeparator:   "$",
+			DisplayID:       false,
+			DisplayName:     true},
+		{Case: "print only id",
+			ExpectedEnabled: true,
+			ExpectedString:  "bar",
+			EnvSubName:      "",
+			EnvSubID:        "",
+			CliExists:       true,
+			CliSubName:      "foo",
+			CliSubID:        "bar",
+			InfoSeparator:   "$",
+			DisplayID:       true,
+			DisplayName:     false},
+		{Case: "print none",
+			ExpectedEnabled: false,
+			ExpectedString:  "",
+			EnvSubName:      "",
+			EnvSubID:        "",
+			CliExists:       true,
+			CliSubName:      "foo",
+			CliSubID:        "bar",
+			InfoSeparator:   "$",
+			DisplayID:       false,
+			DisplayName:     false},
 	}
 
-	a := &az{
+	for _, tc := range cases {
-		env:   env,
+		env := new(MockedEnvironment)
-		props: props,
+		env.On("getenv", "AZ_SUBSCRIPTION_NAME").Return(tc.EnvSubName)
-	}
+		env.On("getenv", "AZ_SUBSCRIPTION_ID").Return(tc.EnvSubID)
-	return a
+		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)
+		props := &properties{
+			values: map[Property]interface{}{
+				SubscriptionInfoSeparator: tc.InfoSeparator,
+				DisplaySubscriptionID:     tc.DisplayID,
+				DisplaySubscriptionName:   tc.DisplayName,
+			},
+		}
 
-func TestEnabledAzNotFound(t *testing.T) {
+		az := &az{
-	args := &azArgs{
+			env:   env,
-		enabled: false,
+			props: props,
+		}
+		assert.Equal(t, tc.ExpectedEnabled, az.enabled(), tc.Case)
+		assert.Equal(t, tc.ExpectedString, az.string(), tc.Case)
 	}
-	az := bootStrapAzTest(args)
-	assert.False(t, az.enabled())
-}
-
-func TestEnabledNoAzDataToDisplay(t *testing.T) {
-	args := &azArgs{
-		enabled:     true,
-		displayID:   false,
-		displayName: false,
-	}
-	az := bootStrapAzTest(args)
-	assert.False(t, az.enabled())
-}
-
-func TestWriteAzSubscriptionId(t *testing.T) {
-	expected := "id"
-	args := &azArgs{
-		enabled:          true,
-		subscriptionID:   "id",
-		subscriptionName: "name",
-		displayID:        true,
-		displayName:      false,
-	}
-	az := bootStrapAzTest(args)
-	assert.True(t, az.enabled())
-	assert.Equal(t, expected, az.string())
-}
-
-func TestWriteAzSubscriptionName(t *testing.T) {
-	expected := "name"
-	args := &azArgs{
-		enabled:          true,
-		subscriptionID:   "id",
-		subscriptionName: "name",
-		displayID:        false,
-		displayName:      true,
-	}
-	az := bootStrapAzTest(args)
-	assert.True(t, az.enabled())
-	assert.Equal(t, expected, az.string())
-}
-
-func TestWriteAzNameAndID(t *testing.T) {
-	expected := "name@id"
-	args := &azArgs{
-		enabled:          true,
-		subscriptionID:   "id",
-		subscriptionName: "name",
-		infoSeparator:    "@",
-		displayID:        true,
-		displayName:      true,
-	}
-	az := bootStrapAzTest(args)
-	assert.True(t, az.enabled())
-	assert.Equal(t, expected, az.string())
 }