Cleanups

Author: javiercn

src/ProjectTemplates/BlazorTemplates.Tests/BlazorServerTemplateTest.cs

@@ -191,8 +191,9 @@ public async Task BlazorServerTemplateWorks_IndividualAuth(BrowserKind browserKi
 
         private async Task TestBasicNavigation(IPage page)
         {
-            IWebSocket socket = null;
-            await page.WaitForEventAsync(PageEvent.WebSocket, s => (socket = s.WebSocket) != null);
+            var socket = BrowserContextInfo.Pages[page].WebSockets.SingleOrDefault() ??
+                (await page.WaitForEventAsync(PageEvent.WebSocket)).WebSocket;
+
             await socket.WaitForEventAsync(WebSocketEvent.FrameReceived);
 
             await page.WaitForSelectorAsync("ul");
@@ -210,6 +211,7 @@ await Task.WhenAll(
 
             // Clicking the counter button works
             await Task.WhenAll(
+                socket.WaitForEventAsync(WebSocketEvent.FrameReceived),
                 page.WaitForSelectorAsync("h1+p >> text=Current count: 1"),
                 page.ClickAsync("p+button >> text=Click me"));
 

src/Shared/BrowserTesting/src/ContextInformation.cs

@@ -34,7 +34,20 @@ private void AttachToPage(object sender, PageEventArgs args)
                 logger.LogInformation($"Network trace will be saved at '{_harPath}'");
             }
 
-            Pages.Add(args.Page, new PageInformation(args.Page, logger));
+            var pageInfo = new PageInformation(args.Page, logger);
+            Pages.Add(args.Page, pageInfo);
+            args.Page.Close += CleanupPage;
+            args.Page.Crash += CleanupPage;
+        }
+
+        private void CleanupPage(object sender, EventArgs e)
+        {
+            var page = (IPage)sender;
+            if (Pages.TryGetValue(page, out var info))
+            {
+                info.Dispose();
+                Pages.Remove(page);
+            }
         }
 
         internal BrowserContextOptions ConfigureUniqueHarPath(BrowserContextOptions browserContextOptions)

src/Shared/BrowserTesting/src/PageInformation.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
@@ -9,7 +9,7 @@
 
 namespace Microsoft.AspNetCore.BrowserTesting
 {
-    public class PageInformation
+    public class PageInformation : IDisposable
     {
         private readonly Page _page;
         private readonly ILogger<PageInformation> _logger;
@@ -20,17 +20,25 @@ public class PageInformation
 
         public List<string> PageErrors { get; } = new();
 
+        public List<IWebSocket> WebSockets { get; set; } = new();
+
         public PageInformation(Page page, ILogger<PageInformation> logger)
         {
             page.Console += RecordConsoleMessage;
             page.PageError += RecordPageError;
             page.RequestFailed += RecordFailedRequest;
+            page.WebSocket += CaptureWebSocket;
             _page = page;
             _logger = logger;
 
             _  = LogPageVideoPath();
         }
 
+        private void CaptureWebSocket(object sender, WebSocketEventArgs e)
+        {
+            WebSockets.Add(e.WebSocket);
+        }
+
         private async Task LogPageVideoPath()
         {
             try
@@ -48,27 +56,58 @@ private async Task LogPageVideoPath()
             }
         }
 
+        public void Dispose()
+        {
+            _page.Console -= RecordConsoleMessage;
+            _page.PageError -= RecordPageError;
+            _page.RequestFailed -= RecordFailedRequest;
+        }
+
         private void RecordFailedRequest(object sender, RequestFailedEventArgs e)
         {
-            _logger.LogError(e.FailureText);
+            try
+            {
+                _logger.LogError(e.FailureText);
+            }
+            catch
+            {
+            }
             FailedRequests.Add(e.FailureText);
         }
 
         private void RecordPageError(object sender, PageErrorEventArgs e)
         {
             // There needs to be a bit of experimentation with this, but message should be a good start.
-            _logger.LogError(e.Message);
+            try
+            {
+                _logger.LogError(e.Message);
+            }
+            catch
+            {
+            }
+
             PageErrors.Add(e.Message);
         }
 
         private void RecordConsoleMessage(object sender, ConsoleEventArgs e)
         {
             var message = e.Message;
+            var messageText = message.Text.Replace(Environment.NewLine, $"{Environment.NewLine}      ");
             var location = message.Location;
 
-            var logMessage = $"[{_page.Url}]{Environment.NewLine}      {message.Text}{Environment.NewLine}      ({location.URL}:{location.LineNumber}:{location.ColumnNumber})";
-            _logger.Log(MapLogLevel(message.Type), logMessage);
-            BrowserConsoleLogs.Add(new LogEntry(message.Text.Replace(Environment.NewLine, $"{Environment.NewLine}      "), message.Type));
+            var logMessage = $"[{_page.Url}]{Environment.NewLine}      {messageText}{Environment.NewLine}      ({location.URL}:{location.LineNumber}:{location.ColumnNumber})";
+
+            try
+            {
+                _logger.Log(MapLogLevel(message.Type), logMessage);
+            }
+            catch
+            {
+
+                throw;
+            }
+
+            BrowserConsoleLogs.Add(new LogEntry(messageText, message.Type));
 
             LogLevel MapLogLevel(string messageType) => messageType switch
             {