log

src/WireMock.Net.Minimal/Matchers/Request/RequestMatchResult.cs

@@ -1,9 +1,5 @@
 // Copyright © WireMock.Net
 
-using System;
-using System.Collections.Generic;
-using System.Linq;
-
 namespace WireMock.Matchers.Request;
 
 /// <summary>

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

@@ -9,6 +9,7 @@ using Stef.Validation;
 using WireMock.Extensions;
 using WireMock.Logging;
 using WireMock.Owin.Mappers;
+using WireMock.Serialization;
 using WireMock.Services;
 using WireMock.Util;
 
@@ -66,6 +67,8 @@ internal partial class AspNetCoreSelfHost
                 services.AddSingleton<IOwinRequestMapper, OwinRequestMapper>();
                 services.AddSingleton<IOwinResponseMapper, OwinResponseMapper>();
                 services.AddSingleton<IGuidUtils, GuidUtils>();
+                services.AddSingleton<LogEntryMapper>();
+                services.AddSingleton<IWireMockMiddlewareLogger, WireMockMiddlewareLogger>();
 
 #if NET8_0_OR_GREATER
                 AddCors(services);

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

@@ -0,0 +1,10 @@
+// Copyright © WireMock.Net
+
+using System.Diagnostics;
+
+namespace WireMock.Owin;
+
+internal interface IWireMockMiddlewareLogger
+{
+    void Log(bool logRequest, RequestMessage request, IResponseMessage? response, MappingMatcherResult? match, MappingMatcherResult? partialMatch, Activity? activity);
+}

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

@@ -1,58 +1,34 @@
 // Copyright © WireMock.Net
 
-using System;
+using System.Diagnostics;
-using System.Collections.Generic;
-using System.Linq;
 using System.Net;
-using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
-using Stef.Validation;
 using WireMock.Constants;
 using WireMock.Exceptions;
 using WireMock.Http;
-using WireMock.Logging;
 using WireMock.Matchers;
+using WireMock.Owin.ActivityTracing;
 using WireMock.Owin.Mappers;
 using WireMock.ResponseBuilders;
 using WireMock.Serialization;
 using WireMock.Settings;
-using WireMock.Util;
-using System.Diagnostics;
-using WireMock.Owin.ActivityTracing;
 
 namespace WireMock.Owin;
 
-internal class WireMockMiddleware
+internal class WireMockMiddleware(
+    RequestDelegate next,
+    IWireMockMiddlewareOptions options,
+    IOwinRequestMapper requestMapper,
+    IOwinResponseMapper responseMapper,
+    IMappingMatcher mappingMatcher,
+    IWireMockMiddlewareLogger logger
+)
 {
     private readonly object _lock = new();
 
-    private readonly IWireMockMiddlewareOptions _options;
-    private readonly IOwinRequestMapper _requestMapper;
-    private readonly IOwinResponseMapper _responseMapper;
-    private readonly IMappingMatcher _mappingMatcher;
-    private readonly LogEntryMapper _logEntryMapper;
-    private readonly IGuidUtils _guidUtils;
-
-    public WireMockMiddleware(
-        RequestDelegate next,
-        IWireMockMiddlewareOptions options,
-        IOwinRequestMapper requestMapper,
-        IOwinResponseMapper responseMapper,
-        IMappingMatcher mappingMatcher,
-        IGuidUtils guidUtils
-    )
-    {
-        _options = Guard.NotNull(options);
-        _requestMapper = Guard.NotNull(requestMapper);
-        _responseMapper = Guard.NotNull(responseMapper);
-        _mappingMatcher = Guard.NotNull(mappingMatcher);
-        _logEntryMapper = new LogEntryMapper(options);
-        _guidUtils = Guard.NotNull(guidUtils);
-    }
-
     public Task Invoke(HttpContext ctx)
     {
-        if (_options.HandleRequestsSynchronously.GetValueOrDefault(false))
+        if (options.HandleRequestsSynchronously.GetValueOrDefault(false))
         {
             lock (_lock)
             {
@@ -66,16 +42,16 @@ internal class WireMockMiddleware
     private async Task InvokeInternalAsync(HttpContext ctx)
     {
         // Store options in HttpContext for providers to access (e.g., WebSocketResponseProvider)
-        ctx.Items[nameof(WireMockMiddlewareOptions)] = _options;
+        ctx.Items[nameof(WireMockMiddlewareOptions)] = options;
 
-        var request = await _requestMapper.MapAsync(ctx, _options).ConfigureAwait(false);
+        var request = await requestMapper.MapAsync(ctx, options).ConfigureAwait(false);
 
         var logRequest = false;
         IResponseMessage? response = null;
         (MappingMatcherResult? Match, MappingMatcherResult? Partial) result = (null, null);
 
-        var tracingEnabled = _options.ActivityTracingOptions is not null;
+        var tracingEnabled = options.ActivityTracingOptions is not null;
-        var excludeAdmin = _options.ActivityTracingOptions?.ExcludeAdminRequests ?? true;
+        var excludeAdmin = options.ActivityTracingOptions?.ExcludeAdminRequests ?? true;
         Activity? activity = null;
 
         // Check if we should trace this request (optionally exclude admin requests)
@@ -84,12 +60,12 @@ internal class WireMockMiddleware
         if (shouldTrace)
         {
             activity = WireMockActivitySource.StartRequestActivity(request.Method, request.Path);
-            WireMockActivitySource.EnrichWithRequest(activity, request, _options.ActivityTracingOptions);
+            WireMockActivitySource.EnrichWithRequest(activity, request, options.ActivityTracingOptions);
         }
 
         try
         {
-            foreach (var mapping in _options.Mappings.Values)
+            foreach (var mapping in options.Mappings.Values)
             {
                 if (mapping.Scenario is null)
                 {
@@ -97,50 +73,50 @@ internal class WireMockMiddleware
                 }
 
                 // Set scenario start
-                if (!_options.Scenarios.ContainsKey(mapping.Scenario) && mapping.IsStartState)
+                if (!options.Scenarios.ContainsKey(mapping.Scenario) && mapping.IsStartState)
                 {
-                    _options.Scenarios.TryAdd(mapping.Scenario, new ScenarioState
+                    options.Scenarios.TryAdd(mapping.Scenario, new ScenarioState
                     {
                         Name = mapping.Scenario
                     });
                 }
             }
 
-            result = _mappingMatcher.FindBestMatch(request);
+            result = mappingMatcher.FindBestMatch(request);
 
             var targetMapping = result.Match?.Mapping;
             if (targetMapping == null)
             {
                 logRequest = true;
-                _options.Logger.Warn("HttpStatusCode set to 404 : No matching mapping found");
+                options.Logger.Warn("HttpStatusCode set to 404 : No matching mapping found");
                 response = ResponseMessageBuilder.Create(HttpStatusCode.NotFound, WireMockConstants.NoMatchingFound);
                 return;
             }
 
             logRequest = targetMapping.LogMapping;
 
-            if (targetMapping.IsAdminInterface && _options.AuthenticationMatcher != null && request.Headers != null)
+            if (targetMapping.IsAdminInterface && options.AuthenticationMatcher != null && request.Headers != null)
             {
                 var authorizationHeaderPresent = request.Headers.TryGetValue(HttpKnownHeaderNames.Authorization, out var authorization);
                 if (!authorizationHeaderPresent)
                 {
-                    _options.Logger.Error("HttpStatusCode set to 401, authorization header is missing.");
+                    options.Logger.Error("HttpStatusCode set to 401, authorization header is missing.");
                     response = ResponseMessageBuilder.Create(HttpStatusCode.Unauthorized, null);
                     return;
                 }
 
-                var authorizationHeaderMatchResult = _options.AuthenticationMatcher.IsMatch(authorization!.ToString());
+                var authorizationHeaderMatchResult = options.AuthenticationMatcher.IsMatch(authorization!.ToString());
                 if (!MatchScores.IsPerfect(authorizationHeaderMatchResult.Score))
                 {
-                    _options.Logger.Error("HttpStatusCode set to 401, authentication failed.", authorizationHeaderMatchResult.Exception ?? throw new WireMockException("Authentication failed"));
+                    options.Logger.Error("HttpStatusCode set to 401, authentication failed.", authorizationHeaderMatchResult.Exception ?? throw new WireMockException("Authentication failed"));
                     response = ResponseMessageBuilder.Create(HttpStatusCode.Unauthorized, null);
                     return;
                 }
             }
 
-            if (!targetMapping.IsAdminInterface && _options.RequestProcessingDelay > TimeSpan.Zero)
+            if (!targetMapping.IsAdminInterface && options.RequestProcessingDelay > TimeSpan.Zero)
             {
-                await Task.Delay(_options.RequestProcessingDelay.Value).ConfigureAwait(false);
+                await Task.Delay(options.RequestProcessingDelay.Value).ConfigureAwait(false);
             }
 
             var (theResponse, theOptionalNewMapping) = await targetMapping.ProvideResponseAsync(ctx, request).ConfigureAwait(false);
@@ -150,7 +126,7 @@ internal class WireMockMiddleware
             {
                 if (responseBuilder?.ProxyAndRecordSettings?.SaveMapping == true || targetMapping.Settings.ProxyAndRecordSettings?.SaveMapping == true)
                 {
-                    _options.Mappings.TryAdd(theOptionalNewMapping.Guid, theOptionalNewMapping);
+                    options.Mappings.TryAdd(theOptionalNewMapping.Guid, theOptionalNewMapping);
                 }
 
                 if (responseBuilder?.ProxyAndRecordSettings?.SaveMappingToFile == true || targetMapping.Settings.ProxyAndRecordSettings?.SaveMappingToFile == true)
@@ -175,56 +151,25 @@ internal class WireMockMiddleware
         }
         catch (Exception ex)
         {
-            _options.Logger.Error($"Providing a Response for Mapping '{result.Match?.Mapping.Guid}' failed. HttpStatusCode set to 500. Exception: {ex}");
+            options.Logger.Error($"Providing a Response for Mapping '{result.Match?.Mapping.Guid}' failed. HttpStatusCode set to 500. Exception: {ex}");
             WireMockActivitySource.RecordException(activity, ex);
 
             response = ResponseMessageBuilder.Create(500, ex.Message);
         }
         finally
         {
-            var log = new LogEntry
+            logger.Log(logRequest, request, response, result.Match, result.Partial, activity);
-            {
-                Guid = _guidUtils.NewGuid(),
-                RequestMessage = request,
-                ResponseMessage = response,
-
-                MappingGuid = result.Match?.Mapping?.Guid,
-                MappingTitle = result.Match?.Mapping?.Title,
-                RequestMatchResult = result.Match?.RequestMatchResult,
-
-                PartialMappingGuid = result.Partial?.Mapping?.Guid,
-                PartialMappingTitle = result.Partial?.Mapping?.Title,
-                PartialMatchResult = result.Partial?.RequestMatchResult
-            };
-
-            WireMockActivitySource.EnrichWithLogEntry(activity, log, _options.ActivityTracingOptions);
-            activity?.Dispose();
-
-            LogRequest(log, logRequest);
 
             try
             {
-                if (_options.SaveUnmatchedRequests == true && result.Match?.RequestMatchResult is not { IsPerfectMatch: true })
+                await responseMapper.MapAsync(response, ctx.Response).ConfigureAwait(false);
-                {
-                    var filename = $"{log.Guid}.LogEntry.json";
-                    _options.FileSystemHandler?.WriteUnmatchedRequest(filename, JsonUtils.Serialize(log));
-                }
-            }
-            catch
-            {
-                // Empty catch
-            }
-
-            try
-            {
-                await _responseMapper.MapAsync(response, ctx.Response).ConfigureAwait(false);
             }
             catch (Exception ex)
             {
-                _options.Logger.Error("HttpStatusCode set to 404 : No matching mapping found", ex);
+                options.Logger.Error("HttpStatusCode set to 404 : No matching mapping found", ex);
 
                 var notFoundResponse = ResponseMessageBuilder.Create(HttpStatusCode.NotFound, WireMockConstants.NoMatchingFound);
-                await _responseMapper.MapAsync(notFoundResponse, ctx.Response).ConfigureAwait(false);
+                await responseMapper.MapAsync(notFoundResponse, ctx.Response).ConfigureAwait(false);
             }
         }
     }
@@ -247,12 +192,12 @@ internal class WireMockMiddleware
                     if (!result.IsSuccessStatusCode)
                     {
                         var content = await result.Content.ReadAsStringAsync().ConfigureAwait(false);
-                        _options.Logger.Warn($"Sending message to Webhook [{webHookIndex}] from Mapping '{mapping.Guid}' failed. HttpStatusCode: {result.StatusCode} Content: {content}");
+                        options.Logger.Warn($"Sending message to Webhook [{webHookIndex}] from Mapping '{mapping.Guid}' failed. HttpStatusCode: {result.StatusCode} Content: {content}");
                     }
                 }
                 catch (Exception ex)
                 {
-                    _options.Logger.Error($"Sending message to Webhook [{webHookIndex}] from Mapping '{mapping.Guid}' failed. Exception: {ex}");
+                    options.Logger.Error($"Sending message to Webhook [{webHookIndex}] from Mapping '{mapping.Guid}' failed. Exception: {ex}");
                 }
             });
         }
@@ -280,7 +225,7 @@ internal class WireMockMiddleware
 
     private void UpdateScenarioState(IMapping mapping)
     {
-        var scenario = _options.Scenarios[mapping.Scenario!];
+        var scenario = options.Scenarios[mapping.Scenario!];
 
         // Increase the number of times this state has been executed
         scenario.Counter++;
@@ -296,59 +241,4 @@ internal class WireMockMiddleware
         scenario.Started = true;
         scenario.Finished = mapping.NextState == null;
     }
-
-    private void LogRequest(LogEntry entry, bool addRequest)
-    {
-        _options.Logger.DebugRequestResponse(_logEntryMapper.Map(entry), entry.RequestMessage.Path.StartsWith("/__admin/"));
-
-        // If addRequest is set to true and MaxRequestLogCount is null or does have a value greater than 0, try to add a new request log.
-        if (addRequest && _options.MaxRequestLogCount is null or > 0)
-        {
-            TryAddLogEntry(entry);
-        }
-
-        // In case MaxRequestLogCount has a value greater than 0, try to delete existing request logs based on the count.
-        if (_options.MaxRequestLogCount is > 0)
-        {
-            var logEntries = _options.LogEntries.ToList();
-            foreach (var logEntry in logEntries.OrderBy(le => le.RequestMessage.DateTime).Take(logEntries.Count - _options.MaxRequestLogCount.Value))
-            {
-                TryRemoveLogEntry(logEntry);
-            }
-        }
-
-        // In case RequestLogExpirationDuration has a value greater than 0, try to delete existing request logs based on the date.
-        if (_options.RequestLogExpirationDuration is > 0)
-        {
-            var checkTime = DateTime.UtcNow.AddHours(-_options.RequestLogExpirationDuration.Value);
-            foreach (var logEntry in _options.LogEntries.ToList().Where(le => le.RequestMessage.DateTime < checkTime))
-            {
-                TryRemoveLogEntry(logEntry);
-            }
-        }
-    }
-
-    private void TryAddLogEntry(LogEntry logEntry)
-    {
-        try
-        {
-            _options.LogEntries.Add(logEntry);
-        }
-        catch
-        {
-            // Ignore exception (can happen during stress testing)
-        }
-    }
-
-    private void TryRemoveLogEntry(LogEntry logEntry)
-    {
-        try
-        {
-            _options.LogEntries.Remove(logEntry);
-        }
-        catch
-        {
-            // Ignore exception (can happen during stress testing)
-        }
-    }
 }

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

@@ -0,0 +1,104 @@
+// Copyright © WireMock.Net
+
+using System.Diagnostics;
+using WireMock.Logging;
+using WireMock.Matchers.Request;
+using WireMock.Owin.ActivityTracing;
+using WireMock.Serialization;
+using WireMock.Util;
+
+namespace WireMock.Owin;
+
+internal class WireMockMiddlewareLogger(IWireMockMiddlewareOptions _options, LogEntryMapper _logEntryMapper, IGuidUtils _guidUtils) : IWireMockMiddlewareLogger
+{
+    public void Log(bool logRequest, RequestMessage request, IResponseMessage? response, MappingMatcherResult? match, MappingMatcherResult? partialMatch, Activity? activity)
+    {
+        var log = new LogEntry
+        {
+            Guid = _guidUtils.NewGuid(),
+            RequestMessage = request,
+            ResponseMessage = response ?? new ResponseMessage(),
+
+            MappingGuid = match?.Mapping?.Guid,
+            MappingTitle = match?.Mapping?.Title,
+            RequestMatchResult = match?.RequestMatchResult ?? new RequestMatchResult(),
+
+            PartialMappingGuid = partialMatch?.Mapping?.Guid,
+            PartialMappingTitle = partialMatch?.Mapping?.Title,
+            PartialMatchResult = partialMatch?.RequestMatchResult ?? new RequestMatchResult()
+        };
+
+        WireMockActivitySource.EnrichWithLogEntry(activity, log, _options.ActivityTracingOptions);
+        activity?.Dispose();
+
+        LogRequest(log, logRequest);
+
+        try
+        {
+            if (_options.SaveUnmatchedRequests == true && match?.RequestMatchResult is not { IsPerfectMatch: true })
+            {
+                var filename = $"{log.Guid}.LogEntry.json";
+                _options.FileSystemHandler?.WriteUnmatchedRequest(filename, JsonUtils.Serialize(log));
+            }
+        }
+        catch
+        {
+            // Empty catch
+        }
+    }
+
+    private void LogRequest(LogEntry entry, bool addRequest)
+    {
+        _options.Logger.DebugRequestResponse(_logEntryMapper.Map(entry), entry.RequestMessage.Path.StartsWith("/__admin/"));
+
+        // If addRequest is set to true and MaxRequestLogCount is null or does have a value greater than 0, try to add a new request log.
+        if (addRequest && _options.MaxRequestLogCount is null or > 0)
+        {
+            TryAddLogEntry(entry);
+        }
+
+        // In case MaxRequestLogCount has a value greater than 0, try to delete existing request logs based on the count.
+        if (_options.MaxRequestLogCount is > 0)
+        {
+            var logEntries = _options.LogEntries.ToList();
+            foreach (var logEntry in logEntries.OrderBy(le => le.RequestMessage.DateTime).Take(logEntries.Count - _options.MaxRequestLogCount.Value))
+            {
+                TryRemoveLogEntry(logEntry);
+            }
+        }
+
+        // In case RequestLogExpirationDuration has a value greater than 0, try to delete existing request logs based on the date.
+        if (_options.RequestLogExpirationDuration is > 0)
+        {
+            var checkTime = DateTime.UtcNow.AddHours(-_options.RequestLogExpirationDuration.Value);
+            foreach (var logEntry in _options.LogEntries.ToList().Where(le => le.RequestMessage.DateTime < checkTime))
+            {
+                TryRemoveLogEntry(logEntry);
+            }
+        }
+    }
+
+    private void TryAddLogEntry(LogEntry logEntry)
+    {
+        try
+        {
+            _options.LogEntries.Add(logEntry);
+        }
+        catch
+        {
+            // Ignore exception (can happen during stress testing)
+        }
+    }
+
+    private void TryRemoveLogEntry(LogEntry logEntry)
+    {
+        try
+        {
+            _options.LogEntries.Remove(logEntry);
+        }
+        catch
+        {
+            // Ignore exception (can happen during stress testing)
+        }
+    }
+}

src/WireMock.Net.Minimal/Serialization/LogEntryMapper.cs

@@ -1,6 +1,5 @@
 // Copyright © WireMock.Net
 
-using System.Linq;
 using Stef.Validation;
 using WireMock.Admin.Mappings;
 using WireMock.Admin.Requests;

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

@@ -250,7 +250,7 @@ internal class WebSocketBuilder(Response response) : IWebSocketBuilder
                 Mapping = context.Mapping,
                 Request = context.RequestMessage,
                 Message = incomingMessage,
-                Data = incomingMessage.MessageType == WebSocketMessageType.Text ? incomingMessage.Text : null
+                Data = context.Mapping.Data
             };
 
             return transformer.Transform(text, model);

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

@@ -2,7 +2,6 @@
 
 using System.Collections.Concurrent;
 using System.Diagnostics.CodeAnalysis;
-using System.Linq;
 using System.Net.WebSockets;
 
 namespace WireMock.WebSockets;

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

@@ -23,7 +23,7 @@ internal struct WebSocketTransformModel
     public WebSocketMessage Message { get; set; }
 
     /// <summary>
-    /// The message data as string
+    /// The mapping data as object
     /// </summary>
-    public string? Data { get; set; }
+    public object? Data { get; set; }
-}
+}

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

@@ -3,7 +3,6 @@
 using System.Net.WebSockets;
 using System.Text;
 using Microsoft.AspNetCore.Http;
-using Newtonsoft.Json;
 using Stef.Validation;
 using WireMock.Extensions;
 using WireMock.Owin;

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

@@ -18,6 +18,8 @@ using WireMock.Matchers.Request;
 using WireMock.ResponseBuilders;
 using WireMock.RequestBuilders;
 using Microsoft.AspNetCore.Http;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
+
 
 #if NET6_0_OR_GREATER
 using WireMock.Owin.ActivityTracing;
@@ -44,8 +46,8 @@ public class WireMockMiddlewareTests
 
     public WireMockMiddlewareTests()
     {
-        var guidUtilsMock = new Mock<IGuidUtils>();
+        var wireMockMiddlewareLoggerMock = new Mock<IWireMockMiddlewareLogger>();
-        guidUtilsMock.Setup(g => g.NewGuid()).Returns(NewGuid);
+        // wreMockMiddlewareLoggerMock.Setup(g => g.NewGuid()).Returns(NewGuid);
 
         _optionsMock = new Mock<IWireMockMiddlewareOptions>();
         _optionsMock.SetupAllProperties();
@@ -84,7 +86,7 @@ public class WireMockMiddlewareTests
             _requestMapperMock.Object,
             _responseMapperMock.Object,
             _matcherMock.Object,
-            guidUtilsMock.Object
+            wireMockMiddlewareLoggerMock.Object
         );
     }
 
@@ -101,28 +103,6 @@ public class WireMockMiddlewareTests
         _responseMapperMock.Verify(m => m.MapAsync(It.Is(match), It.IsAny<HttpResponse>()), Times.Once);
     }
 
-    [Fact]
-    public async Task WireMockMiddleware_Invoke_NoMatch_When_SaveUnmatchedRequestsIsTrue_Should_Call_LocalFileSystemHandler_WriteUnmatchedRequest()
-    {
-        // Arrange
-        var fileSystemHandlerMock = new Mock<IFileSystemHandler>();
-        _optionsMock.Setup(o => o.FileSystemHandler).Returns(fileSystemHandlerMock.Object);
-        _optionsMock.Setup(o => o.SaveUnmatchedRequests).Returns(true);
-
-        // Act
-        await _sut.Invoke(_contextMock.Object);
-
-        // Assert
-        _optionsMock.Verify(o => o.Logger.Warn(It.IsAny<string>(), It.IsAny<object[]>()), Times.Once);
-
-        Expression<Func<ResponseMessage, bool>> match = r => (int)r.StatusCode! == 404 && ((StatusModel)r.BodyData!.BodyAsJson!).Status == "No matching mapping found";
-        _responseMapperMock.Verify(m => m.MapAsync(It.Is(match), It.IsAny<HttpResponse>()), Times.Once);
-
-        // Verify
-        fileSystemHandlerMock.Verify(f => f.WriteUnmatchedRequest("98fae52e-76df-47d9-876f-2ee32e931d9b.LogEntry.json", It.IsAny<string>()));
-        fileSystemHandlerMock.VerifyNoOtherCalls();
-    }
-
     [Fact]
     public async Task WireMockMiddleware_Invoke_IsAdminInterface_EmptyHeaders_401()
     {

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

@@ -11,13 +11,14 @@ using WireMock.Settings;
 
 namespace WireMock.Net.Tests.WebSockets;
 
-public class WebSocketIntegrationTests(ITestOutputHelper output)
+public class WebSocketIntegrationTests(ITestOutputHelper output, ITestContextAccessor testContext)
 {
+    private readonly CancellationToken _ct = testContext.Current.CancellationToken;
+
     [Fact]
     public async Task EchoServer_Should_Echo_Text_Messages()
     {
         // Arrange
-        var cancelationToken = TestContext.Current.CancellationToken;
         using var server = WireMockServer.Start(new WireMockServerSettings
         {
             Logger = new TestOutputHelperWireMockLogger(output),
@@ -43,20 +44,19 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
         client.State.Should().Be(WebSocketState.Open);
 
         var testMessage = "Hello, WebSocket!";
-        await client.SendAsync(testMessage, cancellationToken: cancelationToken);
+        await client.SendAsync(testMessage, cancellationToken: _ct);
 
         // Assert
-        var received = await client.ReceiveAsTextAsync(cancellationToken: cancelationToken);
+        var received = await client.ReceiveAsTextAsync(cancellationToken: _ct);
         received.Should().Be(testMessage);
 
-        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", cancelationToken);
+        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", _ct);
     }
 
     [Fact]
     public async Task WithText_Should_Send_Configured_Text()
     {
         // Arrange
-        var cancelationToken = TestContext.Current.CancellationToken;
         using var server = WireMockServer.Start(new WireMockServerSettings
         {
             Logger = new TestOutputHelperWireMockLogger(output),
@@ -80,24 +80,23 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
         var uri = new Uri($"{server.Url}/ws/message");
 
         // Act
-        await client.ConnectAsync(uri, cancelationToken);
+        await client.ConnectAsync(uri, _ct);
         client.State.Should().Be(WebSocketState.Open);
 
         var testMessage = "Any message from client";
-        await client.SendAsync(testMessage, cancellationToken: cancelationToken);
+        await client.SendAsync(testMessage, cancellationToken: _ct);
 
         // Assert
-        var received = await client.ReceiveAsTextAsync(cancellationToken: cancelationToken);
+        var received = await client.ReceiveAsTextAsync(cancellationToken: _ct);
         received.Should().Be(responseMessage);
 
-        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", cancelationToken);
+        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", _ct);
     }
 
     [Fact]
     public async Task WithText_Should_Send_Same_Text_For_Multiple_Messages()
     {
         // Arrange
-        var cancelationToken = TestContext.Current.CancellationToken;
         using var server = WireMockServer.Start(new WireMockServerSettings
         {
             Logger = new TestOutputHelperWireMockLogger(output),
@@ -119,27 +118,26 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
 
         using var client = new ClientWebSocket();
         var uri = new Uri($"{server.Url}/ws/message");
-        await client.ConnectAsync(uri, cancelationToken);
+        await client.ConnectAsync(uri, _ct);
 
         var testMessages = new[] { "First", "Second", "Third" };
 
         // Act & Assert
         foreach (var testMessage in testMessages)
         {
-            await client.SendAsync(testMessage, cancellationToken: cancelationToken);
+            await client.SendAsync(testMessage, cancellationToken: _ct);
 
-            var received = await client.ReceiveAsTextAsync(cancellationToken: cancelationToken);
+            var received = await client.ReceiveAsTextAsync(cancellationToken: _ct);
             received.Should().Be(responseMessage, $"should always return the fixed response regardless of input message '{testMessage}'");
         }
 
-        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", cancelationToken);
+        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", _ct);
     }
 
     [Fact]
     public async Task WithBinary_Should_Send_Configured_Bytes()
     {
         // Arrange
-        var cancelationToken = TestContext.Current.CancellationToken;
         using var server = WireMockServer.Start(new WireMockServerSettings
         {
             Logger = new TestOutputHelperWireMockLogger(output),
@@ -163,24 +161,23 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
         var uri = new Uri($"{server.Url}/ws/binary");
 
         // Act
-        await client.ConnectAsync(uri, cancelationToken);
+        await client.ConnectAsync(uri, _ct);
         client.State.Should().Be(WebSocketState.Open);
 
         var testMessage = "Any message from client";
-        await client.SendAsync(testMessage, cancellationToken: cancelationToken);
+        await client.SendAsync(testMessage, cancellationToken: _ct);
 
         // Assert
-        var receivedData = await client.ReceiveAsBytesAsync(cancellationToken: cancelationToken);
+        var receivedData = await client.ReceiveAsBytesAsync(cancellationToken: _ct);
         receivedData.Should().BeEquivalentTo(responseBytes);
 
-        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", cancelationToken);
+        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", _ct);
     }
 
     [Fact]
     public async Task WithBinary_Should_Send_Same_Bytes_For_Multiple_Messages()
     {
         // Arrange
-        var cancelationToken = TestContext.Current.CancellationToken;
         using var server = WireMockServer.Start(new WireMockServerSettings
         {
             Logger = new TestOutputHelperWireMockLogger(output),
@@ -202,28 +199,27 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
 
         using var client = new ClientWebSocket();
         var uri = new Uri($"{server.Url}/ws/binary");
-        await client.ConnectAsync(uri, cancelationToken);
+        await client.ConnectAsync(uri, _ct);
 
         var testMessages = new[] { "First", "Second", "Third" };
 
         // Act & Assert
         foreach (var testMessage in testMessages)
         {
-            await client.SendAsync(testMessage, cancellationToken: cancelationToken);
+            await client.SendAsync(testMessage, cancellationToken: _ct);
 
-            var receivedData = await client.ReceiveAsBytesAsync(cancellationToken: cancelationToken);
+            var receivedData = await client.ReceiveAsBytesAsync(cancellationToken: _ct);
             receivedData.Should().BeEquivalentTo(responseBytes, $"should always return the fixed bytes regardless of input message '{testMessage}'");
         }
 
-        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", cancelationToken);
+        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", _ct);
     }
-  
+
 
     [Fact]
     public async Task EchoServer_Should_Echo_Multiple_Messages()
     {
         // Arrange
-        var cancelationToken = TestContext.Current.CancellationToken;
         using var server = WireMockServer.Start(new WireMockServerSettings
         {
             Logger = new TestOutputHelperWireMockLogger(output),
@@ -241,28 +237,27 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
 
         using var client = new ClientWebSocket();
         var uri = new Uri($"{server.Url}/ws/echo");
-        await client.ConnectAsync(uri, cancelationToken);
+        await client.ConnectAsync(uri, _ct);
 
         var testMessages = new[] { "Hello", "World", "WebSocket", "Test" };
 
         // Act & Assert
         foreach (var testMessage in testMessages)
         {
-            await client.SendAsync(testMessage, cancellationToken: cancelationToken);
+            await client.SendAsync(testMessage, cancellationToken: _ct);
 
-            var received = await client.ReceiveAsTextAsync(cancellationToken: cancelationToken);
+            var received = await client.ReceiveAsTextAsync(cancellationToken: _ct);
 
             received.Should().Be(testMessage, $"message '{testMessage}' should be echoed back");
         }
 
-        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", cancelationToken);
+        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", _ct);
     }
 
     [Fact]
     public async Task EchoServer_Should_Echo_Binary_Messages()
     {
         // Arrange
-        var cancelationToken = TestContext.Current.CancellationToken;
         using var server = WireMockServer.Start(new WireMockServerSettings
         {
             Logger = new TestOutputHelperWireMockLogger(output),
@@ -280,26 +275,25 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
 
         using var client = new ClientWebSocket();
         var uri = new Uri($"{server.Url}/ws/echo");
-        await client.ConnectAsync(uri, cancelationToken);
+        await client.ConnectAsync(uri, _ct);
 
         var testData = new byte[] { 0x01, 0x02, 0x03, 0x04, 0x05 };
 
         // Act
-        await client.SendAsync(new ArraySegment<byte>(testData), WebSocketMessageType.Binary, true, cancelationToken);
+        await client.SendAsync(new ArraySegment<byte>(testData), WebSocketMessageType.Binary, true, _ct);
 
-        var receivedData = await client.ReceiveAsBytesAsync(cancellationToken: cancelationToken);
+        var receivedData = await client.ReceiveAsBytesAsync(cancellationToken: _ct);
 
         // Assert
         receivedData.Should().BeEquivalentTo(testData);
 
-        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", cancelationToken);
+        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", _ct);
     }
 
     [Fact]
     public async Task EchoServer_Should_Handle_Empty_Messages()
     {
         // Arrange
-        var cancelationToken = TestContext.Current.CancellationToken;
         using var server = WireMockServer.Start(new WireMockServerSettings
         {
             Logger = new TestOutputHelperWireMockLogger(output),
@@ -317,25 +311,24 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
 
         using var client = new ClientWebSocket();
         var uri = new Uri($"{server.Url}/ws/echo");
-        await client.ConnectAsync(uri, cancelationToken);
+        await client.ConnectAsync(uri, _ct);
 
         // Act
-        await client.SendAsync(string.Empty, cancellationToken: cancelationToken);
+        await client.SendAsync(string.Empty, cancellationToken: _ct);
 
         var receiveBuffer = new byte[1024];
-        var result = await client.ReceiveAsync(new ArraySegment<byte>(receiveBuffer), cancelationToken);
+        var result = await client.ReceiveAsync(new ArraySegment<byte>(receiveBuffer), _ct);
 
         // Assert
         result.Count.Should().Be(0);
 
-        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", cancelationToken);
+        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", _ct);
     }
 
     [Fact]
     public async Task CustomHandler_Should_Handle_Help_Command()
     {
         // Arrange
-        var cancelationToken = TestContext.Current.CancellationToken;
         using var server = WireMockServer.Start(new WireMockServerSettings
         {
             Logger = new TestOutputHelperWireMockLogger(output),
@@ -366,12 +359,12 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
 
         using var client = new ClientWebSocket();
         var uri = new Uri($"{server.Url}/ws/chat");
-        await client.ConnectAsync(uri, cancelationToken);
+        await client.ConnectAsync(uri, _ct);
 
         // Act
-        await client.SendAsync("/help", cancellationToken: cancelationToken);
+        await client.SendAsync("/help", cancellationToken: _ct);
 
-        var received = await client.ReceiveAsTextAsync(cancellationToken: cancelationToken);
+        var received = await client.ReceiveAsTextAsync(cancellationToken: _ct);
 
         // Assert
         received.Should().Contain("Available commands");
@@ -379,14 +372,13 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
         received.Should().Contain("/time");
         received.Should().Contain("/echo");
 
-        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", cancelationToken);
+        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", _ct);
     }
 
     [Fact]
     public async Task CustomHandler_Should_Handle_Multiple_Commands_In_Sequence()
     {
         // Arrange
-        var cancelationToken = TestContext.Current.CancellationToken;
         using var server = WireMockServer.Start(new WireMockServerSettings
         {
             Logger = new TestOutputHelperWireMockLogger(output),
@@ -439,7 +431,7 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
 
         using var client = new ClientWebSocket();
         var uri = new Uri($"{server.Url}/ws/chat");
-        await client.ConnectAsync(uri, cancelationToken);
+        await client.ConnectAsync(uri, _ct);
 
         var commands = new (string, Action<string>)[]
         {
@@ -453,23 +445,22 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
         // Act & Assert
         foreach (var (command, assertion) in commands)
         {
-            await client.SendAsync(command, cancellationToken: cancelationToken);
+            await client.SendAsync(command, cancellationToken: _ct);
 
-            var received = await client.ReceiveAsTextAsync(cancellationToken: cancelationToken);
+            var received = await client.ReceiveAsTextAsync(cancellationToken: _ct);
 
             assertion(received);
         }
 
-        await client.SendAsync("/close", cancellationToken: cancelationToken);
+        await client.SendAsync("/close", cancellationToken: _ct);
 
-        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", cancelationToken);
+        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", _ct);
     }
 
     [Fact]
     public async Task WhenMessage_Should_Handle_Multiple_Conditions_Fluently()
     {
         // Arrange
-        var cancelationToken = TestContext.Current.CancellationToken;
         using var server = WireMockServer.Start(new WireMockServerSettings
         {
             Logger = new TestOutputHelperWireMockLogger(output),
@@ -494,7 +485,7 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
 
         using var client = new ClientWebSocket();
         var uri = new Uri($"{server.Url}/ws/conditional");
-        await client.ConnectAsync(uri, cancelationToken);
+        await client.ConnectAsync(uri, _ct);
 
         var testCases = new (string message, string expectedContains)[]
         {
@@ -508,21 +499,20 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
         // Act & Assert
         foreach (var (message, expectedContains) in testCases)
         {
-            await client.SendAsync(message, cancellationToken: cancelationToken);
+            await client.SendAsync(message, cancellationToken: _ct);
 
-            var received = await client.ReceiveAsTextAsync(cancellationToken: cancelationToken);
+            var received = await client.ReceiveAsTextAsync(cancellationToken: _ct);
 
             received.Should().Contain(expectedContains, $"message '{message}' should return response containing '{expectedContains}'");
         }
 
-        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", cancelationToken);
+        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", _ct);
     }
 
     [Fact]
-    public async Task WhenMessage_NoMatch_Should_Return404()
+    public async Task Request_NoMatch_OnPath_Should_Return404()
     {
         // Arrange
-        var cancelationToken = TestContext.Current.CancellationToken;
         using var server = WireMockServer.Start(new WireMockServerSettings
         {
             Logger = new TestOutputHelperWireMockLogger(output),
@@ -536,48 +526,63 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
             )
             .RespondWith(Response.Create()
                 .WithWebSocket(ws => ws
-                    .WhenMessage("/close")
+                    .WhenMessage("/test")
-                    .ThenSendMessage(m => m.WithText("Closing connection")
+                    .ThenSendMessage(m => m.WithText("Test")
-                    .AndClose()
+                ))
+            );
+
+        using var client = new ClientWebSocket();
+        var uri = new Uri($"{server.Url}/ws/abc");
+
+        // Act
+        Func<Task> connectAction = () => client.ConnectAsync(uri, _ct);
+
+        // Assert
+        (await connectAction.Should().ThrowAsync<WebSocketException>())
+            .WithMessage("The server returned status code '404' when status code '101' was expected.");
+    }
+
+    [Fact]
+    public async Task Request_NoMatch_OnMessageText_Should_ThrowException()
+    {
+        // Arrange
+        using var server = WireMockServer.Start(new WireMockServerSettings
+        {
+            Logger = new TestOutputHelperWireMockLogger(output),
+            Urls = ["ws://localhost:0"]
+        });
+
+        server
+            .Given(Request.Create()
+                .WithPath("/ws/test")
+                .WithWebSocketUpgrade()
+            )
+            .RespondWith(Response.Create()
+                .WithWebSocket(ws => ws
+                    .WithCloseTimeout(TimeSpan.FromSeconds(3))
+                    .WhenMessage("/test")
+                    .ThenSendMessage(m => m.WithText("Test")
                 ))
             );
 
         using var client = new ClientWebSocket();
         var uri = new Uri($"{server.Url}/ws/test");
-        await client.ConnectAsync(uri, cancelationToken);
+
+        await client.ConnectAsync(uri, _ct);
+        await client.SendAsync("/abc", cancellationToken: _ct);
 
         // Act
-        await client.SendAsync("/close", cancellationToken: cancelationToken);
+        Func<Task> receiveAction = () => client.ReceiveAsTextAsync(cancellationToken: _ct);
-
-        var received = await client.ReceiveAsTextAsync(cancellationToken: cancelationToken);
 
         // Assert
-        received.Should().Contain("Closing connection");
+        (await receiveAction.Should().ThrowAsync<WebSocketException>())
-
+            .WithMessage("The remote party closed the WebSocket connection without completing the close handshake.");
-        // Try to receive again - this will complete the close handshake
-        // and update the client state to Closed
-        try
-        {
-            var receiveBuffer = new byte[1024];
-            var result = await client.ReceiveAsync(new ArraySegment<byte>(receiveBuffer), cancelationToken);
-
-            // If we get here, the message type should be Close
-            result.MessageType.Should().Be(WebSocketMessageType.Close);
-        }
-        catch (WebSocketException)
-        {
-            // Connection was closed, which is expected
-        }
-
-        // Verify the connection is CloseReceived
-        client.State.Should().Be(WebSocketState.CloseReceived);
     }
 
     [Fact]
     public async Task WhenMessage_Should_Close_Connection_When_AndClose_Is_Used()
     {
         // Arrange
-        var cancelationToken = TestContext.Current.CancellationToken;
         using var server = WireMockServer.Start(new WireMockServerSettings
         {
             Logger = new TestOutputHelperWireMockLogger(output),
@@ -599,12 +604,12 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
 
         using var client = new ClientWebSocket();
         var uri = new Uri($"{server.Url}/ws/close");
-        await client.ConnectAsync(uri, cancelationToken);
+        await client.ConnectAsync(uri, _ct);
 
         // Act
-        await client.SendAsync("/close", cancellationToken: cancelationToken);
+        await client.SendAsync("/close", cancellationToken: _ct);
 
-        var received = await client.ReceiveAsTextAsync(cancellationToken: cancelationToken);
+        var received = await client.ReceiveAsTextAsync(cancellationToken: _ct);
 
         // Assert
         received.Should().Contain("Closing connection");
@@ -614,8 +619,8 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
         try
         {
             var receiveBuffer = new byte[1024];
-            var result = await client.ReceiveAsync(new ArraySegment<byte>(receiveBuffer), cancelationToken);
+            var result = await client.ReceiveAsync(new ArraySegment<byte>(receiveBuffer), _ct);
-            
+
             // If we get here, the message type should be Close
             result.MessageType.Should().Be(WebSocketMessageType.Close);
         }
@@ -632,7 +637,6 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
     public async Task WithTransformer_Should_Transform_Message_Using_Handlebars()
     {
         // Arrange
-        var cancelationToken = TestContext.Current.CancellationToken;
         using var server = WireMockServer.Start(new WireMockServerSettings
         {
             Logger = new TestOutputHelperWireMockLogger(output),
@@ -655,16 +659,16 @@ public class WebSocketIntegrationTests(ITestOutputHelper output)
         var uri = new Uri($"{server.Url}/ws/transform");
 
         // Act
-        await client.ConnectAsync(uri, cancelationToken);
+        await client.ConnectAsync(uri, _ct);
         client.State.Should().Be(WebSocketState.Open);
 
         var testMessage = "HellO";
-        await client.SendAsync(testMessage, cancellationToken: cancelationToken);
+        await client.SendAsync(testMessage, cancellationToken: _ct);
 
         // Assert
-        var received = await client.ReceiveAsTextAsync(cancellationToken: cancelationToken);
+        var received = await client.ReceiveAsTextAsync(cancellationToken: _ct);
         received.Should().Be("/ws/transform hello");
 
-        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", cancelationToken);
+        await client.CloseAsync(WebSocketCloseStatus.NormalClosure, "Test complete", _ct);
     }
 }