public async Task TestTokenCredentialMultiThreadAsync()
{
// When multiple thread calls TokenCredentialCache.GetTokenAsync and a valid cached token
// is not available, TokenCredentialCache will only create one task to get token.
int numTasks = 100;
TestTokenCredential testTokenCredential = new TestTokenCredential(() =>
{
Task.Delay(TimeSpan.FromSeconds(3)).Wait();
return(new ValueTask <AccessToken>(this.AccessToken));
});
using (TokenCredentialCache tokenCredentialCache = this.CreateTokenCredentialCache(testTokenCredential))
{
Task[] tasks = new Task[numTasks];
for (int i = 0; i < numTasks; i++)
{
tasks[i] = this.GetAndVerifyTokenAsync(tokenCredentialCache);
}
await Task.WhenAll(tasks);
Assert.AreEqual(1, testTokenCredential.NumTimesInvoked);
}
}
public async Task TestTokenCredentialBackgroundRefreshAsync_OnDispose()
{
string token1 = "Token1";
bool firstTimeGetToken = true;
TestTokenCredential testTokenCredential = new TestTokenCredential(() =>
{
if (firstTimeGetToken)
{
firstTimeGetToken = false;
return(new ValueTask <AccessToken>(new AccessToken(token1, DateTimeOffset.UtcNow + TimeSpan.FromSeconds(10))));
}
else
{
throw new Exception("Should not call twice");
}
});
TokenCredentialCache tokenCredentialCache = this.CreateTokenCredentialCache(testTokenCredential, TimeSpan.FromMilliseconds(100));
string t1 = await tokenCredentialCache.GetTokenAsync(NoOpTrace.Singleton);
this.ValidateSemaphoreIsReleased(tokenCredentialCache);
tokenCredentialCache.Dispose();
Assert.AreEqual(token1, t1);
await Task.Delay(1000);
}
public async Task TestTokenCredentialTimeoutAsync()
{
TestTokenCredential testTokenCredential = new TestTokenCredential(async() =>
{
await Task.Delay(-1);
return(new AccessToken("AccessToken", DateTimeOffset.MaxValue));
});
TimeSpan timeout = TimeSpan.FromSeconds(1);
using (TokenCredentialCache tokenCredentialCache = this.CreateTokenCredentialCache(
tokenCredential: testTokenCredential,
requestTimeout: timeout))
{
try
{
await tokenCredentialCache.GetTokenAsync(new CosmosDiagnosticsContextCore());
Assert.Fail("TokenCredentialCache.GetTokenAsync() is expected to fail but succeeded");
}
catch (CosmosException cosmosException)
{
Assert.AreEqual(HttpStatusCode.RequestTimeout, cosmosException.StatusCode);
Assert.AreEqual((int)Azure.Documents.SubStatusCodes.FailedToGetAadToken, cosmosException.SubStatusCode);
Assert.AreEqual($"TokenCredential.GetTokenAsync request timed out after {timeout}", cosmosException.InnerException.Message);
}
}
}
public async Task TestTokenCredentialCacheHappyPathAsync()
{
TestTokenCredential testTokenCredential = new TestTokenCredential(() => new ValueTask <AccessToken>(this.AccessToken));
using (TokenCredentialCache tokenCredentialCache = this.CreateTokenCredentialCache(testTokenCredential))
{
await this.GetAndVerifyTokenAsync(tokenCredentialCache);
}
}
public async Task TestTokenCredentialFailedToRefreshAsync()
{
string token = "Token";
bool firstTimeGetToken = true;
Exception exception = new Exception();
TestTokenCredential testTokenCredential = new TestTokenCredential(() =>
{
if (firstTimeGetToken)
{
firstTimeGetToken = false;
return(new ValueTask <AccessToken>(new AccessToken(token, DateTimeOffset.UtcNow + TimeSpan.FromSeconds(6))));
}
else
{
throw exception;
}
});
using ITrace trace = Trace.GetRootTrace("test");
using (TokenCredentialCache tokenCredentialCache = this.CreateTokenCredentialCache(testTokenCredential))
{
Assert.AreEqual(token, await tokenCredentialCache.GetTokenAsync(trace));
// Token is valid for 6 seconds. Client TokenCredentialRefreshBuffer is set to 5 seconds.
// After waiting for 2 seconds, the cache token is still valid, but it will be refreshed in the background.
await Task.Delay(TimeSpan.FromSeconds(2));
Assert.AreEqual(token, await tokenCredentialCache.GetTokenAsync(trace));
// Token refreshes fails except for the first time, but the cached token will be served as long as it is valid.
await Task.Delay(TimeSpan.FromSeconds(3));
Assert.AreEqual(token, await tokenCredentialCache.GetTokenAsync(trace));
// Cache token has expired, and it fails to refresh.
await Task.Delay(TimeSpan.FromSeconds(2));
try
{
await tokenCredentialCache.GetTokenAsync(trace);
Assert.Fail("TokenCredentialCache.GetTokenAsync() is expected to fail but succeeded");
}
catch (CosmosException cosmosException)
{
Assert.AreEqual(HttpStatusCode.Unauthorized, cosmosException.StatusCode);
Assert.AreEqual((int)Azure.Documents.SubStatusCodes.FailedToGetAadToken, cosmosException.SubStatusCode);
Assert.IsTrue(object.ReferenceEquals(
exception,
cosmosException.InnerException));
}
}
}
public async Task TestTokenCredentialBackgroundRefreshAsync()
{
// When token is within tokenCredentialRefreshBuffer of expiry, start background task to refresh token,
// but return the cached token.
string token1 = "Token1";
string token2 = "Token2";
bool firstTimeGetToken = true;
TestTokenCredential testTokenCredential = new TestTokenCredential(() =>
{
if (firstTimeGetToken)
{
firstTimeGetToken = false;
return(new ValueTask <AccessToken>(new AccessToken(token1, DateTimeOffset.UtcNow + TimeSpan.FromSeconds(10))));
}
else
{
return(new ValueTask <AccessToken>(new AccessToken(token2, DateTimeOffset.MaxValue)));
}
});
using (TokenCredentialCache tokenCredentialCache = this.CreateTokenCredentialCache(testTokenCredential))
{
string t1 = await tokenCredentialCache.GetTokenAsync(NoOpTrace.Singleton);
Assert.AreEqual(token1, t1);
// Token is valid for 6 seconds. Client TokenCredentialRefreshBuffer is set to 5 seconds.
// After waiting for 2 seconds, the cache token is still valid, but it will be refreshed in the background.
await Task.Delay(TimeSpan.FromSeconds(2));
string t2 = await tokenCredentialCache.GetTokenAsync(NoOpTrace.Singleton);
Assert.AreEqual(token1, t2);
// Wait until the background refresh occurs.
while (testTokenCredential.NumTimesInvoked == 1)
{
await Task.Delay(500);
}
string t3 = await tokenCredentialCache.GetTokenAsync(NoOpTrace.Singleton);
Assert.AreEqual(token2, t3);
Assert.AreEqual(2, testTokenCredential.NumTimesInvoked);
this.ValidateSemaphoreIsReleased(tokenCredentialCache);
}
}
public async Task TestTokenCredentialBackgroundRefreshAsync()
{
// When token is within tokenCredentialRefreshBuffer of expiry, start background task to refresh token,
// but return the cached token.
string token1 = "Token1";
string token2 = "Token2";
bool firstTimeGetToken = true;
TestTokenCredential testTokenCredential = new TestTokenCredential(() =>
{
if (firstTimeGetToken)
{
firstTimeGetToken = false;
return(new ValueTask <AccessToken>(new AccessToken(token1, DateTimeOffset.UtcNow + TimeSpan.FromSeconds(10))));
}
else
{
Task.Delay(TimeSpan.FromSeconds(.5)).Wait();
return(new ValueTask <AccessToken>(new AccessToken(token2, DateTimeOffset.MaxValue)));
}
});
using (TokenCredentialCache tokenCredentialCache = this.CreateTokenCredentialCache(testTokenCredential))
{
string t1 = await tokenCredentialCache.GetTokenAsync(new CosmosDiagnosticsContextCore());
Assert.AreEqual(token1, t1);
// Token is valid for 6 seconds. Client TokenCredentialRefreshBuffer is set to 5 seconds.
// After waiting for 2 seconds, the cache token is still valid, but it will be refreshed in the background.
await Task.Delay(TimeSpan.FromSeconds(2));
string t2 = await tokenCredentialCache.GetTokenAsync(new CosmosDiagnosticsContextCore());
Assert.AreEqual(token1, t2);
// After waiting for another 4 seconds (5 seconds for background refresh with .5 second delay), token1 is still valid,
// but cached token has been refreshed to token2 by the background task started before.
await Task.Delay(TimeSpan.FromSeconds(4));
string t3 = await tokenCredentialCache.GetTokenAsync(new CosmosDiagnosticsContextCore());
Assert.AreEqual(token2, t3);
Assert.AreEqual(2, testTokenCredential.NumTimesInvoked);
}
}
public async Task TestTokenCredentialMultiThreadAsync()
{
// When multiple thread calls TokenCredentialCache.GetTokenAsync and a valid cached token
// is not available, TokenCredentialCache will only create one task to get token.
int numTasks = 100;
bool delayTokenRefresh = true;
TestTokenCredential testTokenCredential = new TestTokenCredential(async() =>
{
while (delayTokenRefresh)
{
await Task.Delay(TimeSpan.FromMilliseconds(10));
}
return(this.AccessToken);
});
using (TokenCredentialCache tokenCredentialCache = this.CreateTokenCredentialCache(testTokenCredential))
{
Task[] tasks = new Task[numTasks];
for (int i = 0; i < numTasks; i++)
{
tasks[i] = Task.Run(() => this.GetAndVerifyTokenAsync(tokenCredentialCache));
}
bool waitForTasksToStart = false;
do
{
waitForTasksToStart = tasks.Where(x => x.Status == TaskStatus.Created).Any();
await Task.Delay(TimeSpan.FromMilliseconds(10));
} while (waitForTasksToStart);
// Verify a task took the semaphore lock
bool isRefreshing = false;
do
{
isRefreshing = this.IsTokenRefreshInProgress(tokenCredentialCache);
await Task.Delay(TimeSpan.FromMilliseconds(10));
} while (!isRefreshing);
delayTokenRefresh = false;
await Task.WhenAll(tasks);
this.ValidateSemaphoreIsReleased(tokenCredentialCache);
Assert.AreEqual(1, testTokenCredential.NumTimesInvoked);
}
}
public async Task TestTokenCredentialFailedToRefreshAsync()
{
string token = "Token";
bool throwExceptionOnGetToken = false;
Exception exception = new Exception();
const int semaphoreCount = 10;
using SemaphoreSlim semaphoreSlim = new SemaphoreSlim(semaphoreCount);
TestTokenCredential testTokenCredential = new TestTokenCredential(async() =>
{
try
{
await semaphoreSlim.WaitAsync();
Assert.AreEqual(semaphoreCount - 1, semaphoreSlim.CurrentCount, "Only a single refresh should occur at a time.");
if (throwExceptionOnGetToken)
{
throw exception;
}
else
{
return(new AccessToken(token, DateTimeOffset.UtcNow + TimeSpan.FromSeconds(8)));
}
}
finally
{
semaphoreSlim.Release();
}
});
using ITrace trace = Cosmos.Tracing.Trace.GetRootTrace("test");
using (TokenCredentialCache tokenCredentialCache = this.CreateTokenCredentialCache(testTokenCredential))
{
Assert.AreEqual(token, await tokenCredentialCache.GetTokenAsync(trace));
Assert.AreEqual(1, testTokenCredential.NumTimesInvoked);
throwExceptionOnGetToken = true;
// Token is valid for 10 seconds. Client TokenCredentialRefreshBuffer is set to 5 seconds.
// After waiting for 2 seconds, the cache token is still valid, but it will be refreshed in the background.
await Task.Delay(TimeSpan.FromSeconds(2));
Assert.AreEqual(token, await tokenCredentialCache.GetTokenAsync(trace));
Assert.AreEqual(1, testTokenCredential.NumTimesInvoked);
// Token refreshes fails except for the first time, but the cached token will be served as long as it is valid.
// Wait for the background refresh to occur. It should fail but the cached token should still be valid
Stopwatch stopwatch = Stopwatch.StartNew();
while (testTokenCredential.NumTimesInvoked != 3)
{
Assert.IsTrue(stopwatch.Elapsed.TotalSeconds < 10, "The background task did not start in 10 seconds");
await Task.Delay(200);
}
Assert.AreEqual(token, await tokenCredentialCache.GetTokenAsync(trace));
Assert.AreEqual(3, testTokenCredential.NumTimesInvoked, $"The cached token was not used. Waited time for background refresh: {stopwatch.Elapsed.TotalSeconds} seconds");
// Cache token has expired, and it fails to refresh.
await Task.Delay(TimeSpan.FromSeconds(5));
throwExceptionOnGetToken = true;
// Simulate multiple concurrent request on the failed token
List <Task> tasks = new List <Task>();
for (int i = 0; i < 40; i++)
{
Task task = Task.Run(async() =>
{
try
{
await tokenCredentialCache.GetTokenAsync(trace);
Assert.Fail("TokenCredentialCache.GetTokenAsync() is expected to fail but succeeded");
}
catch (Exception thrownException)
{
// It should just throw the original exception and not be wrapped in a CosmosException
// This avoids any confusion on where the error was thrown from.
Assert.IsTrue(object.ReferenceEquals(
exception,
thrownException), $"Incorrect exception thrown: Expected: {exception}; Actual: {thrownException}");
}
});
tasks.Add(task);
}
await Task.WhenAll(tasks);
this.ValidateSemaphoreIsReleased(tokenCredentialCache);
// Simulate multiple concurrent request that should succeed after a failure
throwExceptionOnGetToken = false;
int numGetTokenCallsAfterFailures = testTokenCredential.NumTimesInvoked;
tasks = new List <Task>();
for (int i = 0; i < 40; i++)
{
Task task = Task.Run(async() => await tokenCredentialCache.GetTokenAsync(trace));
tasks.Add(task);
}
await Task.WhenAll(tasks);
this.ValidateSemaphoreIsReleased(tokenCredentialCache);
}
}
