Merge pull request #2186 from microsoftgraph/enhancements/AddEnvironmentToContext

Add environment to `Get-MgContext` output

Author: peombwa

src/Authentication/Authentication.Core/Interfaces/IAuthContext.cs

@@ -42,6 +42,7 @@ public interface IAuthContext
         string[] Scopes { get; set; }
         AuthenticationType AuthType { get; set; }
         TokenCredentialType TokenCredentialType { get; set; }
+        string Environment { get; set; }
         string AppName { get; set; }
         string Account { get; set; }
         string CertificateThumbprint { get; set; }

src/Authentication/Authentication.Test/Helpers/AuthenticationHelpersTests.cs

@@ -50,6 +50,7 @@ public async Task ShouldUseDelegateAuthProviderWhenUserAccessTokenIsProvidedAsyn
             // Assert
             Assert.IsType<AzureIdentityAccessTokenProvider>(authProvider);
             Assert.Equal(dummyAccessToken, accessToken);
+            Assert.Equal(GraphEnvironmentConstants.EnvironmentName.Global, userProvidedAuthContext.Environment);
 
             // reset static instance.
             GraphSession.Reset();

src/Authentication/Authentication/Cmdlets/ConnectMgGraph.cs

@@ -158,7 +158,7 @@ private async Task ProcessRecordAsync()
         {
             using (NoSynchronizationContext)
             {
-                IAuthContext authContext = new AuthContext { TenantId = TenantId, PSHostVersion = this.Host.Version };
+                IAuthContext authContext = new AuthContext { TenantId = TenantId, PSHostVersion = this.Host.Version, Environment = environment?.Name };
                 if (MyInvocation.BoundParameters.ContainsKey(nameof(ClientTimeout)))
                     GraphSession.Instance.RequestContext.ClientTimeout = TimeSpan.FromSeconds(ClientTimeout);
 

src/Authentication/Authentication/Cmdlets/GetMGContext.cs

@@ -19,7 +19,8 @@ protected override void ProcessRecord()
         {
             base.ProcessRecord();
             IAuthContext authConfig = GraphSession.Instance.AuthContext;
-            WriteObject(authConfig);
+            if (authConfig != null)
+                WriteObject(authConfig);
         }
 
         protected override void EndProcessing()

src/Authentication/Authentication/Models/AuthContext.cs

@@ -2,6 +2,7 @@
 //  Copyright (c) Microsoft Corporation.  All Rights Reserved.  Licensed under the MIT License.  See License in the project root for license information.
 // ------------------------------------------------------------------------------
 
+using Microsoft.Graph.PowerShell.Authentication.Interfaces;
 using System;
 using System.Security;
 using System.Security.Cryptography.X509Certificates;
@@ -25,6 +26,7 @@ public class AuthContext : IAuthContext
         public Version PSHostVersion { get; set; }
         public string ManagedIdentityId { get; set; }
         public SecureString ClientSecret { get; set; }
+        public string Environment { get; set; } = GraphEnvironmentConstants.EnvironmentName.Global;
 
         public AuthContext()
         {

src/Authentication/Authentication/test/Get-MgContext.Tests.ps1

@@ -0,0 +1,39 @@
+# ------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All Rights Reserved. Licensed under the MIT License. See License in the project root for license information.
+# ------------------------------------------------------------------------------
+
+BeforeAll {
+    $ModuleName = "Microsoft.Graph.Authentication"
+    $ModulePath = Join-Path $PSScriptRoot "..\artifacts\$ModuleName.psd1"
+    Import-Module $ModulePath -Force
+    # Test JWT access token. Holds dummy values.
+    $DummyToken = "eyJhbGciOiJIUzI1NiJ9.eyJSb2xlIjoiVGVzdCIsIklzc3VlciI6Iklzc3VlciIsIlVzZXJuYW1lIjoiVGVzdCIsImV4cCI6MTkxMTIzMDM1NiwiaWF0IjoxNjc4NDg4ODE2fQ.yjRvogDyxlQrrQV3EaEsZJKhpYuNzaCyrh5Ip9WvdjU"
+}
+
+Describe 'Get-MgContext' {
+    It 'Should throw return null when no auth context exists for the session' {
+        $Context = Get-MgContext
+        $Context | Should -Be $null
+    }
+    
+    It 'Should default to Global environment when no environment is specified' {
+        Connect-MgGraph -AccessToken (ConvertTo-SecureString -AsPlainText -String $DummyToken)
+        $Context = Get-MgContext
+        $Context | Should -Not -Be $null
+        $Context.Environment | Should -Be "Global"
+        $Context.AuthType | Should -Be "UserProvidedAccessToken"
+    }
+
+        
+    It 'Should return the correct environment when specified via Connect-MgGraph -Environment' {
+        Connect-MgGraph -AccessToken (ConvertTo-SecureString -AsPlainText -String $DummyToken) -Environment Germany
+        $Context = Get-MgContext
+        $Context | Should -Not -Be $null
+        $Context.Environment | Should -Be "Germany"
+        $Context.AuthType | Should -Be "UserProvidedAccessToken"
+    }
+}
+
+AfterAll {
+    Disconnect-MgGraph
+}