diff --git a/test/Microsoft.Azure.SignalR.Tests/ServiceMessageTests.cs b/test/Microsoft.Azure.SignalR.Tests/ServiceMessageTests.cs
index bb5075ba2..a7b996ac2 100644
--- a/test/Microsoft.Azure.SignalR.Tests/ServiceMessageTests.cs
+++ b/test/Microsoft.Azure.SignalR.Tests/ServiceMessageTests.cs
@@ -14,6 +14,7 @@
 using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.SignalR;
 using Microsoft.AspNetCore.SignalR.Internal;
+using Microsoft.Azure.SignalR.Common;
 using Microsoft.Azure.SignalR.Protocol;
 using Microsoft.Azure.SignalR.Tests.Common;
 using Microsoft.Extensions.DependencyInjection;
@@ -243,7 +244,7 @@ public async Task TestAccessKeyResponseMessage(Type keyType)
         {
             case MicrosoftEntraAccessKey key:
                 var source = new CancellationTokenSource(3000);
-                await Assert.ThrowsAsync<TaskCanceledException>(async () => await key.GenerateAccessTokenAsync(DefaultAudience, emptyClaims, lifetime, algorithm, source.Token));
+                await Assert.ThrowsAsync<AzureSignalRAccessTokenNotAuthorizedException>(async () => await key.GenerateAccessTokenAsync(DefaultAudience, emptyClaims, lifetime, algorithm, source.Token));
                 break;
         }