Fixing aadaccesskey test (#1864)

Azure · Oct 26, 2023 · 740cb54 · 740cb54
1 parent b708962
commit 740cb54
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 11 deletions.
diff --git a/src/Microsoft.Azure.SignalR.Common/Endpoints/AadAccessKey.cs b/src/Microsoft.Azure.SignalR.Common/Endpoints/AadAccessKey.cs
@@ -157,7 +157,14 @@ private async Task AuthorizeWithRetryAsync(CancellationToken ctoken = default)
                 catch (Exception e)
                 {
                     latest = e;
-                    await Task.Delay(AuthorizeRetryInterval);
+                    try
+                    {
+                        await Task.Delay(AuthorizeRetryInterval, ctoken);
+                    }
+                    catch (OperationCanceledException)
+                    {
+                        break;
+                    }
                 }
             }
 

diff --git a/test/Microsoft.Azure.SignalR.Common.Tests/Auth/AadAccessKeyTests.cs b/test/Microsoft.Azure.SignalR.Common.Tests/Auth/AadAccessKeyTests.cs
@@ -3,8 +3,9 @@
 using System.Security.Claims;
 using System.Threading;
 using System.Threading.Tasks;
+using Azure.Core;
 using Azure.Identity;
-
+using Moq;
 using Xunit;
 
 namespace Microsoft.Azure.SignalR.Common.Tests.Auth
@@ -27,9 +28,12 @@ public void TestConstructor(string endpoint, string expectedAuthorizeUrl)
         [Fact]
         public async Task TestUpdateAccessKey()
         {
-            var credential = new DefaultAzureCredential();
-            var endpoint = "http://localhost";
-            var key = new AadAccessKey(new Uri(endpoint), credential);
+            var mockCredential = new Mock<TokenCredential>();
+            mockCredential.Setup(credential => credential.GetTokenAsync(
+                It.IsAny<TokenRequestContext>(),
+                It.IsAny<CancellationToken>()))
+                .ThrowsAsync(new InvalidOperationException("Mock GetTokenAsync throws an exception"));
+            var key = new AadAccessKey(new Uri("http://localhost"), mockCredential.Object);
 
             var audience = "http://localhost/chat";
             var claims = Array.Empty<Claim>();
@@ -57,8 +61,12 @@ await Assert.ThrowsAsync<TaskCanceledException>(
         [InlineData(true, 56, false)]
         public async Task TestUpdateAccessKeyShouldSkip(bool isAuthorized, int timeElapsed, bool shouldSkip)
         {
-            var key = new AadAccessKey(new Uri("http://localhost"), new DefaultAzureCredential());
-
+            var mockCredential = new Mock<TokenCredential>();
+            mockCredential.Setup(credential => credential.GetTokenAsync(
+                It.IsAny<TokenRequestContext>(),
+                It.IsAny<CancellationToken>()))
+                .ThrowsAsync(new InvalidOperationException("Mock GetTokenAsync throws an exception"));
+            var key = new AadAccessKey(new Uri("http://localhost"), mockCredential.Object);
             var isAuthorizedField = typeof(AadAccessKey).GetField("_isAuthorized", BindingFlags.NonPublic | BindingFlags.Instance);
             isAuthorizedField.SetValue(key, isAuthorized);
             Assert.Equal(isAuthorized, (bool)isAuthorizedField.GetValue(key));
@@ -81,7 +89,7 @@ public async Task TestUpdateAccessKeyShouldSkip(bool isAuthorized, int timeElaps
             }
             else
             {
-                await Assert.ThrowsAsync<TaskCanceledException>(async () => await key.UpdateAccessKeyAsync(source.Token));
+                await Assert.ThrowsAsync<InvalidOperationException>(async () => await key.UpdateAccessKeyAsync(source.Token));
                 Assert.False((bool)isAuthorizedField.GetValue(key));
                 Assert.True(lastUpdatedTime < (DateTime)lastUpdatedTimeField.GetValue(key));
                 Assert.True(initializedTcs.Task.IsCompleted);
@@ -91,8 +99,12 @@ public async Task TestUpdateAccessKeyShouldSkip(bool isAuthorized, int timeElaps
         [Fact]
         public async Task TestInitializeFailed()
         {
-            var credential = new DefaultAzureCredential();
-            var key = new AadAccessKey(new Uri("http://localhost"), credential);
+            var mockCredential = new Mock<TokenCredential>();
+            mockCredential.Setup(credential => credential.GetTokenAsync(
+                It.IsAny<TokenRequestContext>(),
+                It.IsAny<CancellationToken>()))
+                .ThrowsAsync(new InvalidOperationException("Mock GetTokenAsync throws an exception"));
+            var key = new AadAccessKey(new Uri("http://localhost"), mockCredential.Object);
 
             var audience = "http://localhost/chat";
             var claims = Array.Empty<Claim>();
@@ -104,7 +116,7 @@ public async Task TestInitializeFailed()
             );
 
             var source = new CancellationTokenSource(TimeSpan.FromSeconds(1));
-            await Assert.ThrowsAnyAsync<TaskCanceledException>(
+            await Assert.ThrowsAsync<InvalidOperationException>(
                 async () => await key.UpdateAccessKeyAsync(source.Token)
             );