Improve api for awaiting kernel idle (dotnet#444)

MLS.Agent.Tools/Dotnet.cs

@@ -2,7 +2,6 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System;
-using System.Collections;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;

MLS.Agent/Program.cs

@@ -49,8 +49,11 @@ public static X509Certificate2 ParseKey(string base64EncodedKey)
             typeof(Shell).Assembly
         };
 
-        private static void StartLogging(CompositeDisposable disposables, StartupOptions options)
+        private static IDisposable StartAppInsightsLogging(StartupOptions options)
         {
+
+            var disposables = new CompositeDisposable();
+
             if (options.Production)
             {
                 var applicationVersion = VersionSensor.Version().AssemblyInformationalVersion;
@@ -66,6 +69,24 @@ private static void StartLogging(CompositeDisposable disposables, StartupOptions
                     }));
             }
 
+            if (options.ApplicationInsightsKey != null)
+            {
+                var telemetryClient = new TelemetryClient(new TelemetryConfiguration(options.ApplicationInsightsKey))
+                {
+                    InstrumentationKey = options.ApplicationInsightsKey
+                };
+                disposables.Add(telemetryClient.SubscribeToPocketLogger(_assembliesEmittingPocketLoggerLogs));
+            }
+
+            Log.Event("AgentStarting");
+
+            return disposables;
+        }
+
+        internal static IDisposable StartToolLogging(StartupOptions options)
+        {
+            var disposables = new CompositeDisposable();
+
             if (options.LogPath != null)
             {
                 var log = new SerilogLoggerConfiguration()
@@ -94,22 +115,16 @@ private static void StartLogging(CompositeDisposable disposables, StartupOptions
                 args.SetObserved();
             };
 
-            if (options.ApplicationInsightsKey != null)
-            {
-                var telemetryClient = new TelemetryClient(new TelemetryConfiguration(options.ApplicationInsightsKey))
-                {
-                    InstrumentationKey = options.ApplicationInsightsKey
-                };
-                disposables.Add(telemetryClient.SubscribeToPocketLogger(_assembliesEmittingPocketLoggerLogs));
-            }
-
-            Log.Event("AgentStarting");
+            return disposables;
         }
 
         public static IWebHost ConstructWebHost(StartupOptions options)
         {
-            var disposables = new CompositeDisposable();
-            StartLogging(disposables, options);
+            var disposables = new CompositeDisposable
+            {
+                StartAppInsightsLogging(options),
+                StartToolLogging(options)
+            };
 
             if (options.Key is null)
             {
@@ -120,7 +135,6 @@ public static IWebHost ConstructWebHost(StartupOptions options)
                 Log.Trace("Received Key: {key}", options.Key);
             }
 
-
             var webHost = new WebHostBuilder()
                           .UseKestrel()
                           .UseContentRoot(Path.GetDirectoryName(typeof(Program).Assembly.Location))

Microsoft.DotNet.Interactive.Jupyter.Tests/CompleteRequestHandlerTests.cs

@@ -4,7 +4,9 @@
 using System.Threading.Tasks;
 using Clockwise;
 using FluentAssertions;
+using FluentAssertions.Extensions;
 using Microsoft.DotNet.Interactive.Jupyter.Protocol;
+using Recipes;
 using Xunit;
 using Xunit.Abstractions;
 
@@ -21,9 +23,11 @@ public async Task send_completeReply_on_CompleteRequest()
         {
             var scheduler = CreateScheduler();
             var request = Message.Create(new CompleteRequest("System.Console.", 15), null);
-            await scheduler.Schedule(new JupyterRequestContext(ServerChannel, IoPubChannel, request, KernelStatus));
+            var context = new JupyterRequestContext(ServerChannel, IoPubChannel, request);
 
-            await KernelStatus.Idle();
+            await scheduler.Schedule(context);
+
+            await context.Done().Timeout(5.Seconds());
 
             ServerRecordingSocket.DecodedMessages
                                   .Should()

Microsoft.DotNet.Interactive.Jupyter.Tests/ExecuteRequestHandlerTests.cs

@@ -1,11 +1,12 @@
 // Copyright (c) .NET Foundation and contributors. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
-using System;
 using System.Threading.Tasks;
 using Clockwise;
 using FluentAssertions;
+using FluentAssertions.Extensions;
 using Microsoft.DotNet.Interactive.Jupyter.Protocol;
+using Recipes;
 using Xunit;
 using Xunit.Abstractions;
 
@@ -22,9 +23,10 @@ public async Task sends_ExecuteInput_when_ExecuteRequest_is_handled()
         {
             var scheduler = CreateScheduler();
             var request = Message.Create(new ExecuteRequest("var a =12;"), null);
-            await scheduler.Schedule(new JupyterRequestContext(ServerChannel, IoPubChannel, request, KernelStatus));
+            var context = new JupyterRequestContext(ServerChannel, IoPubChannel, request);
+            await scheduler.Schedule(context);
 
-            await KernelStatus.Idle();
+            await context.Done().Timeout(5.Seconds());
 
             ServerRecordingSocket.DecodedMessages
                                   .Should().Contain(message =>
@@ -40,9 +42,10 @@ public async Task sends_ExecuteReply_message_on_when_code_submission_is_handled(
         {
             var scheduler = CreateScheduler();
             var request = Message.Create(new ExecuteRequest("var a =12;"), null);
-            await scheduler.Schedule(new JupyterRequestContext(ServerChannel, IoPubChannel, request, KernelStatus));
+            var context = new JupyterRequestContext(ServerChannel, IoPubChannel, request);
+            await scheduler.Schedule(context);
 
-            await KernelStatus.Idle();
+            await context.Done().Timeout(5.Seconds());
 
             ServerRecordingSocket.DecodedMessages
                 .Should().Contain(message =>
@@ -54,9 +57,10 @@ public async Task sends_ExecuteReply_with_error_message_on_when_code_submission_
         {
             var scheduler = CreateScheduler();
             var request = Message.Create(new ExecuteRequest("asdes"), null);
-            await scheduler.Schedule(new JupyterRequestContext(ServerChannel, IoPubChannel, request, KernelStatus));
+            var context = new JupyterRequestContext(ServerChannel, IoPubChannel, request);
+            await scheduler.Schedule(context);
 
-            await KernelStatus.Idle();
+            await context.Done().Timeout(5.Seconds());
 
             ServerRecordingSocket.DecodedMessages
                 .Should()
@@ -74,9 +78,10 @@ public async Task sends_DisplayData_message_on_ValueProduced()
         {
             var scheduler = CreateScheduler();
             var request = Message.Create(new ExecuteRequest("Console.WriteLine(2+2);"), null);
-            await scheduler.Schedule(new JupyterRequestContext(ServerChannel, IoPubChannel, request, KernelStatus));
+            var context = new JupyterRequestContext(ServerChannel, IoPubChannel, request);
+            await scheduler.Schedule(context);
 
-            await KernelStatus.Idle();
+            await context.Done().Timeout(10.Seconds());
 
             ServerRecordingSocket.DecodedMessages
                 .Should().Contain(message =>
@@ -92,9 +97,10 @@ public async Task sends_ExecuteReply_message_on_ReturnValueProduced()
         {
             var scheduler = CreateScheduler();
             var request = Message.Create(new ExecuteRequest("2+2"), null);
-            await scheduler.Schedule(new JupyterRequestContext(ServerChannel, IoPubChannel, request, KernelStatus));
+            var context = new JupyterRequestContext(ServerChannel, IoPubChannel, request);
+            await scheduler.Schedule(context);
 
-            await KernelStatus.Idle();
+            await context.Done().Timeout(5.Seconds());
 
             ServerRecordingSocket.DecodedMessages
                 .Should().Contain(message =>
@@ -110,9 +116,10 @@ public async Task sends_ExecuteReply_message_when_submission_contains_only_a_dir
         {
             var scheduler = CreateScheduler();
             var request = Message.Create(new ExecuteRequest("%%csharp"), null);
-            await scheduler.Schedule(new JupyterRequestContext(ServerChannel, IoPubChannel, request, KernelStatus));
+            var context = new JupyterRequestContext(ServerChannel, IoPubChannel, request);
+            await scheduler.Schedule(context);
 
-            await KernelStatus.Idle();
+            await context.Done().Timeout(5.Seconds());
 
             ServerRecordingSocket.DecodedMessages
                 .Should().Contain(message =>

Microsoft.DotNet.Interactive.Jupyter.Tests/InterruptRequestHandlerTests.cs

@@ -1,12 +1,13 @@
 // Copyright (c) .NET Foundation and contributors. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
-using System;
 using Clockwise;
 using FluentAssertions;
 using System.Linq;
 using System.Threading.Tasks;
+using FluentAssertions.Extensions;
 using Microsoft.DotNet.Interactive.Jupyter.Protocol;
+using Recipes;
 using Xunit;
 using Xunit.Abstractions;
 
@@ -23,9 +24,11 @@ public async Task sends_InterruptReply()
         {
             var scheduler = CreateScheduler();
             var request = Message.Create(new InterruptRequest(), null);
-            await scheduler.Schedule(new JupyterRequestContext(ServerChannel, IoPubChannel, request, KernelStatus));
+            var context = new JupyterRequestContext(ServerChannel, IoPubChannel, request);
 
-            await KernelStatus.Idle();
+            await scheduler.Schedule(context);
+
+            await context.Done().Timeout(5.Seconds());
 
             ServerRecordingSocket.DecodedMessages
                                   .SingleOrDefault(message =>

Microsoft.DotNet.Interactive.Jupyter.Tests/IsCompleteRequestHandlerTests.cs

@@ -5,6 +5,7 @@
 using FluentAssertions;
 using System.Linq;
 using System.Threading.Tasks;
+using FluentAssertions.Extensions;
 using Microsoft.DotNet.Interactive.Jupyter.Protocol;
 using Pocket;
 using Recipes;
@@ -24,9 +25,10 @@ public async Task sends_isCompleteReply_with_complete_if_the_code_is_a_complete_
         {
             var scheduler = CreateScheduler();
             var request = Message.Create(new IsCompleteRequest("var a = 12;"), null);
+            var context = new JupyterRequestContext(ServerChannel, IoPubChannel, request);
 
-            await scheduler.Schedule(new JupyterRequestContext(ServerChannel, IoPubChannel, request, KernelStatus));
-            await KernelStatus.Idle();
+            await scheduler.Schedule(context);
+            await context.Done().Timeout(5.Seconds());
 
             Logger.Log.Info("DecodedMessages: {messages}", ServerRecordingSocket.DecodedMessages);
 
@@ -46,8 +48,11 @@ public async Task sends_isCompleteReply_with_incomplete_and_indent_if_the_code_i
         {
             var scheduler = CreateScheduler();
             var request = Message.Create(new IsCompleteRequest("var a = 12"), null);
-            await scheduler.Schedule(new JupyterRequestContext(ServerChannel, IoPubChannel, request, KernelStatus));
-            await KernelStatus.Idle();
+            var context = new JupyterRequestContext(ServerChannel, IoPubChannel, request);
+
+            await scheduler.Schedule(context);
+            await context.Done().Timeout(5.Seconds());
+
             ServerRecordingSocket.DecodedMessages.SingleOrDefault(message =>
                                                                        message.Contains(JupyterMessageContentTypes.IsCompleteReply))
                                   .Should()

Microsoft.DotNet.Interactive.Jupyter.Tests/JupyterRequestHandlerTestBase{T}.cs

@@ -13,35 +13,36 @@ namespace Microsoft.DotNet.Interactive.Jupyter.Tests
     public abstract class JupyterRequestHandlerTestBase<T> : IDisposable
         where T : JupyterMessageContent
     {
-        private readonly CompositeDisposable _disposables =new CompositeDisposable();
-        protected MessageSender IoPubChannel { get; }
-        protected MessageSender ServerChannel { get; }
-        protected RecordingSocket ServerRecordingSocket { get; }
-        protected RecordingSocket IoRecordingSocket { get; }
-        protected KernelStatus KernelStatus { get; }
+        private readonly CompositeDisposable _disposables = new CompositeDisposable();
+        protected readonly MessageSender IoPubChannel;
+        protected readonly MessageSender ServerChannel;
+        protected readonly RecordingSocket ServerRecordingSocket;
+        protected readonly RecordingSocket IoRecordingSocket;
+        protected readonly IKernel Kernel;
 
         protected JupyterRequestHandlerTestBase(ITestOutputHelper output)
         {
             _disposables.Add(output.SubscribeToPocketLogger());
 
+            Kernel = new CompositeKernel
+            {
+                new CSharpKernel()
+            }.UseDefaultMagicCommands();
+
+            _disposables.Add(Kernel.LogEventsToPocketLogger());
+
             var signatureValidator = new SignatureValidator("key", "HMACSHA256");
             ServerRecordingSocket = new RecordingSocket();
             ServerChannel = new MessageSender(ServerRecordingSocket, signatureValidator);
             IoRecordingSocket = new RecordingSocket();
             IoPubChannel = new MessageSender(IoRecordingSocket, signatureValidator);
-            KernelStatus = new KernelStatus(
-                Header.Create<T>("test"),
-                ServerChannel);
         }
 
         public void Dispose() => _disposables.Dispose();
 
         protected ICommandScheduler<JupyterRequestContext> CreateScheduler()
         {
-            var handler = new JupyterRequestContextHandler(new CompositeKernel
-            {
-                new CSharpKernel()
-            }.UseDefaultMagicCommands());
+            var handler = new JupyterRequestContextHandler(Kernel);
 
             return CommandScheduler.Create<JupyterRequestContext>(handler.Handle).Trace();
         }