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 = {
@@ -47,7 +58,7 @@ function global:Set-PoshGitStatus {
         $executionTime = ($history.EndExecutionTime - $history.StartExecutionTime).TotalMilliseconds
         $global:omp_lastHistoryId = $history.Id
     }
-    Set-PoshGitStatus
+    Initialize-ModuleSupport
     $omp = "::OMP::"
     $config = $global:PoshSettings.Theme
     $cleanPWD = $PWD.ProviderPath.TrimEnd("\")

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
+		CliSubID        string
+		InfoSeparator   string
+		DisplayID       bool
+		DisplayName     bool
+	}{
+		{Case: "envvars present",
+			ExpectedEnabled: true,
+			ExpectedString:  "foo$bar",
+			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},
 	}
 
-func bootStrapAzTest(args *azArgs) *az {
+	for _, tc := range cases {
 		env := new(MockedEnvironment)
-	env.On("hasCommand", "az").Return(args.enabled)
+		env.On("getenv", "AZ_SUBSCRIPTION_NAME").Return(tc.EnvSubName)
-	env.On("runCommand", "az", []string{"account", "show", "--query=[name,id]", "-o=tsv"}).Return(fmt.Sprintf("%s\n%s\n", args.subscriptionName, args.subscriptionID), nil)
+		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)
 		props := &properties{
 			values: map[Property]interface{}{
-			SubscriptionInfoSeparator: args.infoSeparator,
+				SubscriptionInfoSeparator: tc.InfoSeparator,
-			DisplaySubscriptionID:     args.displayID,
+				DisplaySubscriptionID:     tc.DisplayID,
-			DisplaySubscriptionName:   args.displayName,
+				DisplaySubscriptionName:   tc.DisplayName,
 			},
 		}
 
-	a := &az{
+		az := &az{
 			env:   env,
 			props: props,
 		}
-	return a
+		assert.Equal(t, tc.ExpectedEnabled, az.enabled(), tc.Case)
+		assert.Equal(t, tc.ExpectedString, az.string(), tc.Case)
 	}
-
-func TestEnabledAzNotFound(t *testing.T) {
-	args := &azArgs{
-		enabled: false,
-	}
-	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())
 }