[dotnet] [bidi] Make BiDi transport factories composable

dotnet/src/webdriver/BiDi/BiDi.cs

@@ -60,7 +60,11 @@ public static async Task<IBiDi> ConnectAsync(Uri url, Action<BiDiOptionsBuilder>
         BiDiOptionsBuilder builder = new();
         configure?.Invoke(builder);
 
-        var transport = await builder.TransportFactory(url, cancellationToken).ConfigureAwait(false);
+        var transportFactoryTask = builder.TransportFactory(url, cancellationToken)
+            ?? throw new InvalidOperationException("The transport factory must return a non-null Task<ITransport> instance.");
+
+        var transport = await transportFactoryTask.ConfigureAwait(false)
+            ?? throw new InvalidOperationException("The transport factory must return a non-null ITransport instance.");
 
         BiDi bidi = new();
 

dotnet/src/webdriver/BiDi/BiDiOptionsBuilder.cs

@@ -27,8 +27,11 @@ namespace OpenQA.Selenium.BiDi;
 /// </summary>
 public sealed class BiDiOptionsBuilder
 {
+    private static readonly Func<Uri, CancellationToken, Task<ITransport>> DefaultTransportFactory =
+        (uri, ct) => WebSocketTransport.ConnectAsync(uri, null, ct);
+
     internal Func<Uri, CancellationToken, Task<ITransport>> TransportFactory { get; private set; }
-        = (uri, ct) => WebSocketTransport.ConnectAsync(uri, null, ct);
+        = DefaultTransportFactory;
 
     /// <summary>
     /// Configures the BiDi connection to use a WebSocket transport.
@@ -42,36 +45,27 @@ public sealed class BiDiOptionsBuilder
     /// <returns>The current <see cref="BiDiOptionsBuilder"/> instance for chaining.</returns>
     public BiDiOptionsBuilder UseWebSocket(Action<ClientWebSocketOptions>? configure = null)
     {
-        return UseTransport((uri, ct) => WebSocketTransport.ConnectAsync(uri, configure, ct));
+        TransportFactory = (uri, ct) => WebSocketTransport.ConnectAsync(uri, configure, ct);
+        return this;
     }
 
     /// <summary>
-    /// Configures the BiDi connection to use a transport created by the specified factory.
+    /// Composes a transport factory into the current transport pipeline.
     /// </summary>
     /// <remarks>
-    /// BiDi takes ownership of the transport instance returned by the factory and will dispose it.
+    /// The <paramref name="next"/> callback receives the current transport factory and returns
+    /// the next factory in the chain. BiDi takes ownership of the transport instance returned by
+    /// the final factory and will dispose it.
     /// </remarks>
-    /// <param name="factory">A factory function that creates the <see cref="ITransport"/> instance.</param>
+    /// <param name="next">A callback that composes a new transport factory from the current one.</param>
     /// <returns>The current <see cref="BiDiOptionsBuilder"/> instance for chaining.</returns>
-    public BiDiOptionsBuilder UseTransport(Func<ITransport> factory)
+    public BiDiOptionsBuilder UseTransport(Func<Func<Uri, CancellationToken, Task<ITransport>>, Func<Uri, CancellationToken, Task<ITransport>>> next)
     {
-        ArgumentNullException.ThrowIfNull(factory);
-
-        return UseTransport((_, ct) =>
-        {
-            if (ct.IsCancellationRequested)
-            {
-                return Task.FromCanceled<ITransport>(ct);
-            }
+        ArgumentNullException.ThrowIfNull(next);
 
-            var transport = factory() ?? throw new InvalidOperationException("The transport factory must return a non-null ITransport instance.");
+        var factory = next(TransportFactory)
+            ?? throw new InvalidOperationException("The transport factory decorator must return a non-null factory.");
 
-            return Task.FromResult(transport);
-        });
-    }
-
-    private BiDiOptionsBuilder UseTransport(Func<Uri, CancellationToken, Task<ITransport>> factory)
-    {
         TransportFactory = factory;
         return this;
     }

dotnet/test/webdriver/BiDi/SessionUnitTests.cs

@@ -34,7 +34,8 @@ class SessionUnitTests
     public async Task SetUp()
     {
         _transport = new FakeTransport();
-        _bidi = await Selenium.BiDi.BiDi.ConnectAsync(new Uri("ws://fake"), opts => opts.UseTransport(() => _transport));
+        _bidi = await Selenium.BiDi.BiDi.ConnectAsync(new Uri("ws://fake"), opts =>
+            opts.UseTransport(_ => (_, _) => Task.FromResult<ITransport>(_transport)));
     }
 
     [TearDown]