From 7d5b1a92ee1111315b0aa832874a3dcbf571842f Mon Sep 17 00:00:00 2001
From: Vincent Dai <23257217+vidai-msft@users.noreply.github.com>
Date: Thu, 14 Dec 2023 15:50:02 +0800
Subject: [PATCH] Fix graph client url issue in testfx

---
 .../ScenarioTests/ResourcesTestRunner.cs              |  2 +-
 tools/TestFx/Mocks/MockContext.cs                     | 11 ++++++-----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/Resources/Resources.Test/ScenarioTests/ResourcesTestRunner.cs b/src/Resources/Resources.Test/ScenarioTests/ResourcesTestRunner.cs
index 0f1a0b7acd48..f6827fb95670 100644
--- a/src/Resources/Resources.Test/ScenarioTests/ResourcesTestRunner.cs
+++ b/src/Resources/Resources.Test/ScenarioTests/ResourcesTestRunner.cs
@@ -62,7 +62,7 @@ protected ResourcesTestRunner(ITestOutputHelper output)
                         ? new Func<ServiceClientCredentials>(() =>
                         {
                             var testEnvironment = TestEnvironmentFactory.GetTestEnvironment();
-                            return testEnvironment.TokenInfo[TokenAudience.Graph];
+                            return testEnvironment.TokenInfo[TokenAudience.Management];
                         })()
                         : new TokenCredentials("foo");
 
diff --git a/tools/TestFx/Mocks/MockContext.cs b/tools/TestFx/Mocks/MockContext.cs
index d1784d72a333..e244b782e2e1 100644
--- a/tools/TestFx/Mocks/MockContext.cs
+++ b/tools/TestFx/Mocks/MockContext.cs
@@ -106,7 +106,7 @@ public T GetServiceClient<T>(TestEnvironment currentEnvironment, bool internalBa
                     $"Unable to create Service Client because {nameof(T)} authentication token was not acquired during Login.");
             }
 
-            return GetServiceClientWithCredentials<T>(currentEnvironment, currentEnvironment.TokenInfo[TokenAudience.Management], internalBaseUri, handlers);
+            return GetServiceClientWithCredentials<T>(currentEnvironment, currentEnvironment.BaseUri, currentEnvironment.TokenInfo[TokenAudience.Management], internalBaseUri, handlers);
         }
 
         /// <summary>
@@ -135,7 +135,7 @@ public T GetGraphServiceClient<T>(TestEnvironment currentEnvironment, bool inter
                     "Unable to create Graph Client because Graph authentication token was not acquired during Login.");
             }
 
-            return GetServiceClientWithCredentials<T>(currentEnvironment, currentEnvironment.TokenInfo[TokenAudience.Graph], internalBaseUri, handlers);
+            return GetServiceClientWithCredentials<T>(currentEnvironment, currentEnvironment.Endpoints.GraphUri, currentEnvironment.TokenInfo[TokenAudience.Graph], internalBaseUri, handlers);
         }
 
         /// <summary>
@@ -146,14 +146,15 @@ public T GetGraphServiceClient<T>(TestEnvironment currentEnvironment, bool inter
         /// <param name="baseUri">Base Uri</param>
         /// <param name="handlers">Delegating existingHandlers</param>
         /// <returns></returns>
-        public T GetServiceClientWithCredentials<T>(TestEnvironment currentEnvironment, object credentials, bool internalBaseUri = false, params DelegatingHandler[] handlers) where T : class
+        public T GetServiceClientWithCredentials<T>(TestEnvironment currentEnvironment, Uri baseUri, object credentials, bool internalBaseUri, params DelegatingHandler[] handlers) where T : class
         {
             T client;
             handlers = AddHandlers(currentEnvironment, handlers);
             var constructors = typeof(T).GetConstructors(BindingFlags.Public | BindingFlags.Static | BindingFlags.Instance | BindingFlags.NonPublic);
 
+            Uri uri = baseUri ?? currentEnvironment.BaseUri;
             ConstructorInfo constructor = null;
-            if (!internalBaseUri && currentEnvironment.BaseUri != null)
+            if (!internalBaseUri && uri != null)
             {
                 foreach (var c in constructors)
                 {
@@ -173,7 +174,7 @@ public T GetServiceClientWithCredentials<T>(TestEnvironment currentEnvironment,
                 }
                 client = constructor.Invoke(new object[]
                 {
-                    currentEnvironment.BaseUri,
+                    uri,
                     credentials,
                     handlers
                 }) as T;