public async Task OnceSessionReleasedGetSessionWillQueueWorkOnIt() { var livyClient = Substitute.For <ILivyClient>(); var livySession1 = Substitute.For <ILivySession>(); var livySession2 = Substitute.For <ILivySession>(); var livySession3 = Substitute.For <ILivySession>(); var livySession4 = Substitute.For <ILivySession>(); var livySession5 = Substitute.For <ILivySession>(); var source1 = new TaskCompletionSource <bool>(); var source2 = new TaskCompletionSource <bool>(); var source3 = new TaskCompletionSource <bool>(); var source4 = new TaskCompletionSource <bool>(); var source5 = new TaskCompletionSource <bool>(); livySession1.WaitForSessionAsync().Returns(source1.Task); livySession2.WaitForSessionAsync().Returns(source2.Task); livySession3.WaitForSessionAsync().Returns(source3.Task); livySession4.WaitForSessionAsync().Returns(source4.Task); livySession5.WaitForSessionAsync().Returns(source5.Task); livyClient.CreateSessionAsync(Arg.Any <LivySessionConfiguration>()) .Returns(Task.FromResult(livySession1), Task.FromResult(livySession2), Task.FromResult(livySession3), Task.FromResult(livySession4), Task.FromResult(livySession5), null); var pool = new CosmosDbLivySessionPool(livyClient, new CosmosCollectionSettings(), LivySessionConfiguration.Default, 5); async void WaitAndReleaseSession3() { await Task.Delay(100); source3.SetResult(true); } WaitAndReleaseSession3(); var result = await Enumerable.Repeat(0, 100).Select(_ => pool.GetSessionAsync()); var sessions = result.GroupBy(s => s).Select(g => new { g.Key, Count = g.Count() }); ILivySession GetSession(CosmosDbLivySession sqlSession) => ((Lazy <Task <ILivySession> >) typeof(CosmosDbLivySession).GetField("_session", BindingFlags.NonPublic | BindingFlags.Instance)? .GetValue(sqlSession))?.Value.Result; var enumerable = sessions.Single(g => GetSession(g.Key) == livySession3); Assert.That(enumerable.Count, Is.EqualTo(96)); }
public async Task GetSessionAlwaysReturnOneOfFiveSessions() { var livyClient = Substitute.For <ILivyClient>(); var livySession1 = Substitute.For <ILivySession>(); var livySession2 = Substitute.For <ILivySession>(); var livySession3 = Substitute.For <ILivySession>(); var livySession4 = Substitute.For <ILivySession>(); var livySession5 = Substitute.For <ILivySession>(); livyClient.CreateSessionAsync(Arg.Any <LivySessionConfiguration>()) .Returns(Task.FromResult(livySession1), Task.FromResult(livySession2), Task.FromResult(livySession3), Task.FromResult(livySession4), Task.FromResult(livySession5), null); var pool = new CosmosDbLivySessionPool(livyClient, new CosmosCollectionSettings(), LivySessionConfiguration.Default, 5); var result = await Enumerable.Repeat(0, 100).Select(_ => pool.GetSessionAsync()); Assert.That(result.Distinct().Count(), Is.EqualTo(5)); }