Reuse session HTTP client. (#1184)

Author: peombwa

src/Authentication/Authentication.Core/Common/GraphSession.cs

@@ -9,6 +9,7 @@ namespace Microsoft.Graph.PowerShell.Authentication
 
     using System;
     using System.Collections;
+    using System.Net.Http;
     using System.Security;
     using System.Threading;
 
@@ -68,6 +69,8 @@ public byte[] MSALToken
         /// </summary>
         public Hashtable[] MgCommandMetadata { get; set; }
 
+        public HttpClient GraphHttpClient { get; set; }
+
         /// <summary>
         /// Gets an instance of <see cref="GraphSession"/>.
         /// </summary>

src/Authentication/Authentication.Test/Helpers/HttpHelpersTests.cs

@@ -5,6 +5,7 @@ namespace Microsoft.Graph.Authentication.Test.Helpers
 {
     using System;
     using System.Net.Http;
+    using Microsoft.Graph.Authentication.Core;
     using Microsoft.Graph.PowerShell.Authentication;
     using Microsoft.Graph.PowerShell.Authentication.Helpers;
     using Xunit;
@@ -71,5 +72,112 @@ public void GetGraphHttpClientShouldReturnHttpClientWithCustomerProvidedTimeout(
             // reset static instance.
             GraphSession.Reset();
         }
+
+        [Fact]
+        public void GetGraphHttpClientShouldReturnInSessionHttpClientWhenSessionHasAClient()
+        {
+            GraphSession.Initialize(() => new GraphSession());
+            GraphSession.Instance.AuthContext = new AuthContext
+            {
+                AuthType = AuthenticationType.UserProvidedAccessToken,
+                ContextScope = ContextScope.Process,
+            };
+            GraphSession.Instance.GraphHttpClient = new HttpClient
+            {
+                BaseAddress = new Uri("https://test.contoso.com/v1.0/")
+            };
+
+            HttpClient httpClient = HttpHelpers.GetGraphHttpClient();
+
+            Assert.NotNull(httpClient);
+            Assert.NotNull(GraphSession.Instance.GraphHttpClient);
+            Assert.Equal(httpClient.BaseAddress, GraphSession.Instance.GraphHttpClient.BaseAddress);
+
+            // reset static instance.
+            GraphSession.Reset();
+        }
+
+        [Fact]
+        public void GetGraphHttpClientShouldReturnNewHttpClientWhenSessionHasNoClient()
+        {
+            GraphSession.Initialize(() => new GraphSession());
+            GraphSession.Instance.AuthContext = new AuthContext
+            {
+                AuthType = AuthenticationType.UserProvidedAccessToken,
+                ContextScope = ContextScope.Process,
+            };
+            GraphSession.Instance.GraphHttpClient = null;
+
+            HttpClient httpClient = HttpHelpers.GetGraphHttpClient();
+
+            Assert.NotNull(httpClient);
+            Assert.NotNull(GraphSession.Instance.GraphHttpClient);
+            Assert.Equal(httpClient.BaseAddress, GraphSession.Instance.GraphHttpClient.BaseAddress);
+
+            // reset static instance.
+            GraphSession.Reset();
+        }
+
+        [Fact]
+        public void GetGraphHttpClientShouldReturnNewHttpClientWhenSessionIsNew()
+        {
+            GraphSession.Initialize(() => new GraphSession());
+            GraphSession.Instance.AuthContext = new AuthContext
+            {
+                AuthType = AuthenticationType.UserProvidedAccessToken,
+                ContextScope = ContextScope.Process,
+            };
+
+            HttpClient httpClient = HttpHelpers.GetGraphHttpClient();
+
+            Assert.NotNull(httpClient);
+            Assert.NotNull(GraphSession.Instance.GraphHttpClient);
+            Assert.Equal(httpClient.BaseAddress, GraphSession.Instance.GraphHttpClient.BaseAddress);
+
+            // reset static instance.
+            GraphSession.Reset();
+        }
+
+        [Fact]
+        public void GetGraphHttpClientShouldReturnNewHttpClientSignOutThenSignIn()
+        {
+            GraphSession.Initialize(() => new GraphSession());
+            GraphSession.Instance.AuthContext = new AuthContext
+            {
+                AuthType = AuthenticationType.UserProvidedAccessToken,
+                ContextScope = ContextScope.Process,
+            };
+
+            var dummyClient = new HttpClient
+            {
+                BaseAddress = new Uri("https://test.contoso.com/v1.0/")
+            };
+            GraphSession.Instance.GraphHttpClient = dummyClient;
+            HttpClient httpClientAttempt1 = HttpHelpers.GetGraphHttpClient();
+
+            // Mock sign out.
+            Authenticator.LogOut(GraphSession.Instance.AuthContext);
+            GraphSession.Instance.AuthContext = null;
+            GraphSession.Instance.GraphHttpClient = null;
+
+            // Mock sign in.
+            GraphSession.Initialize(() => new GraphSession());
+            GraphSession.Instance.AuthContext = new AuthContext
+            {
+                AuthType = AuthenticationType.UserProvidedAccessToken,
+                ContextScope = ContextScope.Process,
+            };
+            HttpClient httpClientAttempt2 = HttpHelpers.GetGraphHttpClient();
+
+            Assert.NotNull(httpClientAttempt1);
+            Assert.NotNull(httpClientAttempt2);
+            Assert.NotNull(GraphSession.Instance.GraphHttpClient);
+            Assert.NotEqual(httpClientAttempt2.BaseAddress, httpClientAttempt1.BaseAddress);
+            Assert.Equal(httpClientAttempt1.BaseAddress, dummyClient.BaseAddress);
+            Assert.Equal(httpClientAttempt2.BaseAddress, GraphSession.Instance.GraphHttpClient.BaseAddress);
+
+            // reset static instance.
+            GraphSession.Reset();
+        }
     }
 }

src/Authentication/Authentication/Cmdlets/ConnectMgGraph.cs

@@ -189,6 +189,7 @@ private async Task ProcessRecordAsync()
                 if (MyInvocation.BoundParameters.ContainsKey(nameof(ClientTimeout))) { authContext.ClientTimeout = TimeSpan.FromSeconds(ClientTimeout); }
                 // Set selected environment to the session object.
                 GraphSession.Instance.Environment = environment;
+                GraphSession.Instance.GraphHttpClient = null;
                 switch (ParameterSetName)
                 {
                     case Constants.UserParameterSet:

src/Authentication/Authentication/Cmdlets/DisconnectMgGraph.cs

@@ -34,6 +34,7 @@ protected override void ProcessRecord()
             Authenticator.LogOut(authContext);
 
             GraphSession.Instance.AuthContext = null;
+            GraphSession.Instance.GraphHttpClient = null;
         }
 
         protected override void StopProcessing()

src/Authentication/Authentication/Helpers/HttpHelpers.cs

@@ -60,14 +60,18 @@ public static HttpClient GetGraphHttpClient(InvocationInfo invocationInfo, IAuth
         /// <returns></returns>
         public static HttpClient GetGraphHttpClient(IAuthContext authContext = null)
         {
+            if (GraphSession.Instance?.GraphHttpClient != null)
+                return GraphSession.Instance.GraphHttpClient;
             authContext = authContext ?? GraphSession.Instance.AuthContext;
             if (authContext is null)
             {
                 throw new AuthenticationException(Core.ErrorConstants.Message.MissingAuthContext);
             }
 
             IAuthenticationProvider authProvider = AuthenticationHelpers.GetAuthProvider(authContext);
-            return GetGraphHttpClient(authProvider, authContext.ClientTimeout);
+            var newHttpClient = GetGraphHttpClient(authProvider, authContext.ClientTimeout);
+            GraphSession.Instance.GraphHttpClient = newHttpClient;
+            return newHttpClient;
         }
 
         /// <summary>