From 90e29360a9eeb82d5463511345f8885a8aaa8d51 Mon Sep 17 00:00:00 2001
From: Michael Render <render.michael@gmail.com>
Date: Fri, 14 Mar 2025 15:11:54 -0400
Subject: [PATCH] [dotnet] [bidi] Encapsulate transport inside `Broker`

---
 dotnet/src/webdriver/BiDi/BiDi.cs                 | 13 +++++++------
 dotnet/src/webdriver/BiDi/Communication/Broker.cs | 14 +++++++++++---
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/dotnet/src/webdriver/BiDi/BiDi.cs b/dotnet/src/webdriver/BiDi/BiDi.cs
index 1ec5138156a6b..1feb6ebddfa1a 100644
--- a/dotnet/src/webdriver/BiDi/BiDi.cs
+++ b/dotnet/src/webdriver/BiDi/BiDi.cs
@@ -20,13 +20,11 @@
 using System;
 using System.Threading.Tasks;
 using OpenQA.Selenium.BiDi.Communication;
-using OpenQA.Selenium.BiDi.Communication.Transport;
 
 namespace OpenQA.Selenium.BiDi;
 
 public class BiDi : IAsyncDisposable
 {
-    private readonly ITransport _transport;
     private readonly Broker _broker;
 
     private readonly Lazy<Modules.Session.SessionModule> _sessionModule;
@@ -42,8 +40,7 @@ internal BiDi(string url)
     {
         var uri = new Uri(url);
 
-        _transport = new WebSocketTransport(new Uri(url));
-        _broker = new Broker(this, _transport);
+        _broker = new Broker(this, uri);
 
         _sessionModule = new Lazy<Modules.Session.SessionModule>(() => new Modules.Session.SessionModule(_broker));
         _browsingContextModule = new Lazy<Modules.BrowsingContext.BrowsingContextModule>(() => new Modules.BrowsingContext.BrowsingContextModule(_broker));
@@ -83,10 +80,14 @@ public Task EndAsync(Modules.Session.EndOptions? options = null)
         return SessionModule.EndAsync(options);
     }
 
-    public async ValueTask DisposeAsync()
+    public virtual async ValueTask DisposeAsyncCore()
     {
         await _broker.DisposeAsync().ConfigureAwait(false);
+    }
 
-        _transport?.Dispose();
+    public async ValueTask DisposeAsync()
+    {
+        await DisposeAsyncCore();
+        GC.SuppressFinalize(this);
     }
 }
diff --git a/dotnet/src/webdriver/BiDi/Communication/Broker.cs b/dotnet/src/webdriver/BiDi/Communication/Broker.cs
index de7fd33a35180..bd999ed947e13 100644
--- a/dotnet/src/webdriver/BiDi/Communication/Broker.cs
+++ b/dotnet/src/webdriver/BiDi/Communication/Broker.cs
@@ -54,10 +54,10 @@ public class Broker : IAsyncDisposable
 
     private readonly BiDiJsonSerializerContext _jsonSerializerContext;
 
-    internal Broker(BiDi bidi, ITransport transport)
+    internal Broker(BiDi bidi, Uri url)
     {
         _bidi = bidi;
-        _transport = transport;
+        _transport = new WebSocketTransport(url);
 
         var jsonSerializerOptions = new JsonSerializerOptions
         {
@@ -298,7 +298,7 @@ public async Task UnsubscribeAsync(Modules.Session.Subscription subscription, Ev
         }
     }
 
-    public async ValueTask DisposeAsync()
+    public virtual async ValueTask DisposeAsyncCore()
     {
         _pendingEvents.CompleteAdding();
 
@@ -308,5 +308,13 @@ public async ValueTask DisposeAsync()
         {
             await _eventEmitterTask.ConfigureAwait(false);
         }
+
+        _transport.Dispose();
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        await DisposeAsyncCore();
+        GC.SuppressFinalize(this);
     }
 }