diff --git a/dotnet/src/webdriver/BiDi/BiDi.cs b/dotnet/src/webdriver/BiDi/BiDi.cs index 93dd627890567..3a51763890ba0 100644 --- a/dotnet/src/webdriver/BiDi/BiDi.cs +++ b/dotnet/src/webdriver/BiDi/BiDi.cs @@ -60,7 +60,11 @@ public static async Task ConnectAsync(Uri url, Action 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 instance."); + + var transport = await transportFactoryTask.ConfigureAwait(false) + ?? throw new InvalidOperationException("The transport factory must return a non-null ITransport instance."); BiDi bidi = new(); diff --git a/dotnet/src/webdriver/BiDi/BiDiOptionsBuilder.cs b/dotnet/src/webdriver/BiDi/BiDiOptionsBuilder.cs index 23b7d433fde4f..06c279cec02a7 100644 --- a/dotnet/src/webdriver/BiDi/BiDiOptionsBuilder.cs +++ b/dotnet/src/webdriver/BiDi/BiDiOptionsBuilder.cs @@ -27,8 +27,11 @@ namespace OpenQA.Selenium.BiDi; /// public sealed class BiDiOptionsBuilder { + private static readonly Func> DefaultTransportFactory = + (uri, ct) => WebSocketTransport.ConnectAsync(uri, null, ct); + internal Func> TransportFactory { get; private set; } - = (uri, ct) => WebSocketTransport.ConnectAsync(uri, null, ct); + = DefaultTransportFactory; /// /// Configures the BiDi connection to use a WebSocket transport. @@ -42,36 +45,27 @@ public sealed class BiDiOptionsBuilder /// The current instance for chaining. public BiDiOptionsBuilder UseWebSocket(Action? configure = null) { - return UseTransport((uri, ct) => WebSocketTransport.ConnectAsync(uri, configure, ct)); + TransportFactory = (uri, ct) => WebSocketTransport.ConnectAsync(uri, configure, ct); + return this; } /// - /// Configures the BiDi connection to use a transport created by the specified factory. + /// Composes a transport factory into the current transport pipeline. /// /// - /// BiDi takes ownership of the transport instance returned by the factory and will dispose it. + /// The 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. /// - /// A factory function that creates the instance. + /// A callback that composes a new transport factory from the current one. /// The current instance for chaining. - public BiDiOptionsBuilder UseTransport(Func factory) + public BiDiOptionsBuilder UseTransport(Func>, Func>> next) { - ArgumentNullException.ThrowIfNull(factory); - - return UseTransport((_, ct) => - { - if (ct.IsCancellationRequested) - { - return Task.FromCanceled(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> factory) - { TransportFactory = factory; return this; } diff --git a/dotnet/test/webdriver/BiDi/SessionUnitTests.cs b/dotnet/test/webdriver/BiDi/SessionUnitTests.cs index 8afa9eb9a9c9c..cd699cdc5adb7 100644 --- a/dotnet/test/webdriver/BiDi/SessionUnitTests.cs +++ b/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(_transport))); } [TearDown]