...

src/WireMock.Net.Minimal/Owin/Mappers/IOwinRequestMapper.cs

@@ -1,12 +1,6 @@
 // Copyright © WireMock.Net
 
-using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
-//#if !USE_ASPNETCORE
-//using IRequest = Microsoft.Owin.IOwinRequest;
-//#else
-//using IRequest = Microsoft.AspNetCore.Http.HttpRequest;
-//#endif
 
 namespace WireMock.Owin.Mappers;
 
@@ -18,8 +12,8 @@ internal interface IOwinRequestMapper
     /// <summary>
     /// MapAsync IRequest to RequestMessage
     /// </summary>
-    /// <param name="request">The HttpRequest</param>
+    /// <param name="context">The HttpContext</param>
     /// <param name="options">The WireMockMiddlewareOptions</param>
     /// <returns>RequestMessage</returns>
-    Task<RequestMessage> MapAsync(HttpRequest request, IWireMockMiddlewareOptions options);
+    Task<RequestMessage> MapAsync(HttpContext context, IWireMockMiddlewareOptions options);
 }

src/WireMock.Net.Minimal/Owin/Mappers/OwinRequestMapper.cs

@@ -1,9 +1,5 @@
 // Copyright © WireMock.Net
 
-using System;
-using System.Collections.Generic;
-using System.Linq;
-using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Extensions;
 using WireMock.Http;
@@ -18,8 +14,9 @@ namespace WireMock.Owin.Mappers;
 internal class OwinRequestMapper : IOwinRequestMapper
 {
     /// <inheritdoc />
-    public async Task<RequestMessage> MapAsync(HttpRequest request, IWireMockMiddlewareOptions options)
+    public async Task<RequestMessage> MapAsync(HttpContext context, IWireMockMiddlewareOptions options)
     {
+        var request = context.Request;
         var (urlDetails, clientIP) = ParseRequest(request);
 
         var method = request.Method;

src/WireMock.Net.Minimal/Owin/WireMockMiddleware.cs

@@ -68,7 +68,7 @@ internal class WireMockMiddleware
         // Store options in HttpContext for providers to access (e.g., WebSocketResponseProvider)
         ctx.Items[nameof(WireMockMiddlewareOptions)] = _options;
 
-        var request = await _requestMapper.MapAsync(ctx.Request, _options).ConfigureAwait(false);
+        var request = await _requestMapper.MapAsync(ctx, _options).ConfigureAwait(false);
 
         var logRequest = false;
         IResponseMessage? response = null;

src/WireMock.Net.Minimal/RequestBuilders/Request.WithWebSocket.cs

@@ -1,6 +1,5 @@
 // Copyright © WireMock.Net
 
-using System.Linq;
 using WireMock.Matchers;
 using WireMock.Matchers.Request;
 
@@ -8,6 +7,9 @@ namespace WireMock.RequestBuilders;
 
 public partial class Request
 {
+    /// <inheritdoc />
+    public bool IsWebSocket { get; private set; }
+
     /// <inheritdoc />
     public IRequestBuilder WithWebSocketUpgrade(params string[] protocols)
     {
@@ -38,6 +40,8 @@ public partial class Request
             ));
         }
 
+        IsWebSocket = true;
+
         return this;
     }
 }

src/WireMock.Net.Minimal/RequestBuilders/Request.cs

@@ -2,11 +2,8 @@
 
 // This source file is based on mock4net by Alexandre Victoor which is licensed under the Apache 2.0 License.
 // For more details see 'mock4net/LICENSE.txt' and 'mock4net/readme.md' in this project root.
-using System;
-using System.Collections.Generic;
 using System.Collections.ObjectModel;
 using System.Diagnostics.CodeAnalysis;
-using System.Linq;
 using Stef.Validation;
 using WireMock.Matchers;
 using WireMock.Matchers.Request;

src/WireMock.Net.Minimal/ResponseProviders/WebSocketResponseProvider.cs

@@ -1,11 +1,8 @@
 // Copyright © WireMock.Net
 
-using System;
 using System.Net;
 using System.Net.WebSockets;
 using System.Text;
-using System.Threading;
-using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Stef.Validation;
 using WireMock.Constants;
@@ -58,7 +55,7 @@ internal class WebSocketResponseProvider : IResponseProvider
                 throw new InvalidOperationException("WireMockMiddlewareOptions not found in HttpContext.Items");
             }
 
-            // Get or create registry from options (not from server)
+            // Get or create registry from options
             var registry = _builder.IsBroadcast 
                 ? options.WebSocketRegistries.GetOrAdd(mapping.Guid, _ => new WebSocketConnectionRegistry())
                 : null;
@@ -113,10 +110,7 @@ internal class WebSocketResponseProvider : IResponseProvider
             finally
             {
                 // Remove from registry
-                if (registry != null)
+                registry?.RemoveConnection(wsContext.ConnectionId);
-                {
-                    registry.RemoveConnection(wsContext.ConnectionId);
-                }
             }
 
             // Return special marker to indicate WebSocket was handled

src/WireMock.Net.Minimal/WebSockets/WebSocketConnectionRegistry.cs

@@ -1,13 +1,8 @@
 // Copyright © WireMock.Net
 
-using System;
 using System.Collections.Concurrent;
-using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
-using System.Linq;
 using System.Net.WebSockets;
-using System.Threading;
-using System.Threading.Tasks;
 using Newtonsoft.Json;
 
 namespace WireMock.WebSockets;

src/WireMock.Net.Shared/RequestBuilders/IWebSocketRequestBuilder.cs

@@ -8,6 +8,11 @@ namespace WireMock.RequestBuilders;
 /// </summary>
 public interface IWebSocketRequestBuilder : IRequestMatcher
 {
+    /// <summary>
+    /// Gets a value indicating whether the connection uses the WebSocket protocol.
+    /// </summary>
+    bool IsWebSocket { get; }
+
     /// <summary>
     /// Match WebSocket upgrade with optional protocols.
     /// </summary>

test/WireMock.Net.Tests/Owin/WireMockMiddlewareTests.cs

@@ -72,7 +72,7 @@ public class WireMockMiddlewareTests
         _requestMapperMock = new Mock<IOwinRequestMapper>();
         _requestMapperMock.SetupAllProperties();
         var request = new RequestMessage(new UrlDetails("http://localhost/foo"), "GET", "::1");
-        _requestMapperMock.Setup(m => m.MapAsync(It.IsAny<IRequest>(), It.IsAny<IWireMockMiddlewareOptions>())).ReturnsAsync(request);
+        _requestMapperMock.Setup(m => m.MapAsync(It.IsAny<HttpContext>(), It.IsAny<IWireMockMiddlewareOptions>())).ReturnsAsync(request);
 
         _responseMapperMock = new Mock<IOwinResponseMapper>();
         _responseMapperMock.SetupAllProperties();
@@ -141,7 +141,7 @@ public class WireMockMiddlewareTests
     {
         // Assign
         var request = new RequestMessage(new UrlDetails("http://localhost/foo"), "GET", "::1", null, new Dictionary<string, string[]>());
-        _requestMapperMock.Setup(m => m.MapAsync(It.IsAny<IRequest>(), It.IsAny<IWireMockMiddlewareOptions>())).ReturnsAsync(request);
+        _requestMapperMock.Setup(m => m.MapAsync(It.IsAny<HttpContext>(), It.IsAny<IWireMockMiddlewareOptions>())).ReturnsAsync(request);
 
         _optionsMock.SetupGet(o => o.AuthenticationMatcher).Returns(new ExactMatcher());
         _mappingMock.SetupGet(m => m.IsAdminInterface).Returns(true);
@@ -164,7 +164,7 @@ public class WireMockMiddlewareTests
     {
         // Assign
         var request = new RequestMessage(new UrlDetails("http://localhost/foo"), "GET", "::1", null, new Dictionary<string, string[]> { { "h", new[] { "x" } } });
-        _requestMapperMock.Setup(m => m.MapAsync(It.IsAny<IRequest>(), It.IsAny<IWireMockMiddlewareOptions>())).ReturnsAsync(request);
+        _requestMapperMock.Setup(m => m.MapAsync(It.IsAny<HttpContext>(), It.IsAny<IWireMockMiddlewareOptions>())).ReturnsAsync(request);
 
         _optionsMock.SetupGet(o => o.AuthenticationMatcher).Returns(new ExactMatcher());
         _mappingMock.SetupGet(m => m.IsAdminInterface).Returns(true);
@@ -197,7 +197,7 @@ public class WireMockMiddlewareTests
     {
         // Assign
         var request = new RequestMessage(new UrlDetails("http://localhost/foo"), "GET", "::1", null, new Dictionary<string, string[]>());
-        _requestMapperMock.Setup(m => m.MapAsync(It.IsAny<IRequest>(), It.IsAny<IWireMockMiddlewareOptions>())).ReturnsAsync(request);
+        _requestMapperMock.Setup(m => m.MapAsync(It.IsAny<HttpContext>(), It.IsAny<IWireMockMiddlewareOptions>())).ReturnsAsync(request);
 
         _optionsMock.SetupGet(o => o.AuthenticationMatcher).Returns(new ExactMatcher());
 
@@ -246,7 +246,7 @@ public class WireMockMiddlewareTests
     {
         // Assign
         var request = new RequestMessage(new UrlDetails("http://localhost/foo"), "GET", "::1", null, new Dictionary<string, string[]>());
-        _requestMapperMock.Setup(m => m.MapAsync(It.IsAny<IRequest>(), It.IsAny<IWireMockMiddlewareOptions>())).ReturnsAsync(request);
+        _requestMapperMock.Setup(m => m.MapAsync(It.IsAny<HttpContext>(), It.IsAny<IWireMockMiddlewareOptions>())).ReturnsAsync(request);
 
         _optionsMock.SetupGet(o => o.AuthenticationMatcher).Returns(new ExactMatcher());
 
@@ -301,7 +301,7 @@ public class WireMockMiddlewareTests
     {
         // Arrange
         var request = new RequestMessage(new UrlDetails("http://localhost/__admin/health"), "GET", "::1");
-        _requestMapperMock.Setup(m => m.MapAsync(It.IsAny<IRequest>(), It.IsAny<IWireMockMiddlewareOptions>())).ReturnsAsync(request);
+        _requestMapperMock.Setup(m => m.MapAsync(It.IsAny<HttpContext>(), It.IsAny<IWireMockMiddlewareOptions>())).ReturnsAsync(request);
 
         _optionsMock.SetupGet(o => o.ActivityTracingOptions).Returns(new ActivityTracingOptions
         {
@@ -330,7 +330,7 @@ public class WireMockMiddlewareTests
     {
         // Arrange
         var request = new RequestMessage(new UrlDetails("http://localhost/api/orders"), "GET", "::1");
-        _requestMapperMock.Setup(m => m.MapAsync(It.IsAny<IRequest>(), It.IsAny<IWireMockMiddlewareOptions>())).ReturnsAsync(request);
+        _requestMapperMock.Setup(m => m.MapAsync(It.IsAny<HttpContext>(), It.IsAny<IWireMockMiddlewareOptions>())).ReturnsAsync(request);
 
         _optionsMock.SetupGet(o => o.ActivityTracingOptions).Returns(new ActivityTracingOptions
         {
@@ -359,7 +359,7 @@ public class WireMockMiddlewareTests
     {
         // Arrange
         var request = new RequestMessage(new UrlDetails("http://localhost/api/orders"), "GET", "::1");
-        _requestMapperMock.Setup(m => m.MapAsync(It.IsAny<IRequest>(), It.IsAny<IWireMockMiddlewareOptions>())).ReturnsAsync(request);
+        _requestMapperMock.Setup(m => m.MapAsync(It.IsAny<HttpContext>(), It.IsAny<IWireMockMiddlewareOptions>())).ReturnsAsync(request);
 
         _optionsMock.SetupGet(o => o.ActivityTracingOptions).Returns((ActivityTracingOptions?)null);
 

test/WireMock.Net.Tests/WebSockets/WebSocketIntegrationTests.cs

@@ -35,6 +35,7 @@ public class WebSocketIntegrationTests
         server
             .Given(Request.Create()
                 .WithPath("/ws/echo")
+                //.WithBody("Hello, WebSocket!")
                 .WithWebSocketUpgrade()
             )
             .RespondWith(Response.Create()