BodyType

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

@@ -7,8 +7,6 @@ using System.Text;
 using Microsoft.AspNetCore.Http;
 using Stef.Validation;
 using WireMock.Logging;
-using WireMock.Matchers;
-using WireMock.Matchers.Request;
 using WireMock.Models;
 using WireMock.Owin;
 using WireMock.Owin.ActivityTracing;
@@ -93,51 +91,97 @@ public class WireMockWebSocketContext : IWebSocketContext
         );
     }
 
-    private async Task SendAsyncInternal(
-        ArraySegment<byte> buffer,
-        WebSocketMessageType messageType,
-        bool endOfMessage,
-        object? data,
-        CancellationToken cancellationToken)
+    /// <inheritdoc />
+    public async Task CloseAsync(WebSocketCloseStatus closeStatus, string statusDescription)
     {
-        Activity? activity = null;
-        var shouldTrace = Options.ActivityTracingOptions is not null;
+        await WebSocket.CloseAsync(closeStatus, statusDescription, CancellationToken.None);
 
-        if (shouldTrace)
+        LogWebSocketMessage(WebSocketMessageDirection.Send, WebSocketMessageType.Close, $"CloseStatus: {closeStatus}, Description: {statusDescription}", null);
+    }
+
+    /// <inheritdoc />
+    public void SetScenarioState(string nextState)
+    {
+        SetScenarioState(nextState, null);
+    }
+
+    /// <inheritdoc />
+    public void SetScenarioState(string nextState, string? description)
+    {
+        if (Mapping.Scenario == null)
         {
-            activity = WireMockActivitySource.StartWebSocketMessageActivity(WebSocketMessageDirection.Send, Mapping.Guid);
-            WireMockActivitySource.EnrichWithWebSocketMessage(
-                activity,
-                messageType,
-                buffer.Count,
-                endOfMessage,
-                data as string,
-                Options.ActivityTracingOptions
-            );
+            return;
         }
 
-        try
+        // Use the same logic as WireMockMiddleware
+        if (Options.Scenarios.TryGetValue(Mapping.Scenario, out var scenarioState))
         {
-            await WebSocket.SendAsync(buffer, messageType, endOfMessage, cancellationToken).ConfigureAwait(false);
+            // Directly set the next state (bypass counter logic for manual WebSocket state changes)
+            scenarioState.NextState = nextState;
+            scenarioState.Started = true;
+            scenarioState.Finished = nextState == null;
 
-            // Log the send operation
-            if (Options.MaxRequestLogCount is null or > 0)
+            // Reset counter when manually setting state
+            scenarioState.Counter = 0;
+        }
+        else
+        {
+            // Create new scenario state if it doesn't exist
+            Options.Scenarios.TryAdd(Mapping.Scenario, new ScenarioState
             {
-                LogWebSocketMessage(WebSocketMessageDirection.Send, messageType, data, activity);
-            }
-        }
-        catch (Exception ex)
-        {
-            WireMockActivitySource.RecordException(activity, ex);
-            throw;
-        }
-        finally
-        {
-            activity?.Dispose();
+                Name = Mapping.Scenario,
+                NextState = nextState,
+                Started = true,
+                Finished = nextState == null,
+                Counter = 0
+            });
         }
     }
 
-    private void LogWebSocketMessage(
+    /// <inheritdoc />
+    public async Task BroadcastTextAsync(string text, CancellationToken cancellationToken = default)
+    {
+        if (Registry != null)
+        {
+            await Registry.BroadcastTextAsync(text, cancellationToken);
+        }
+    }
+
+    /// <summary>
+    /// Update scenario state following the same pattern as WireMockMiddleware.UpdateScenarioState
+    /// This is called automatically when the WebSocket connection is established.
+    /// </summary>
+    internal void UpdateScenarioState()
+    {
+        if (Mapping.Scenario == null)
+        {
+            return;
+        }
+
+        // Ensure scenario exists
+        if (!Options.Scenarios.TryGetValue(Mapping.Scenario, out var scenario))
+        {
+            return;
+        }
+
+        // Follow exact same logic as WireMockMiddleware.UpdateScenarioState
+        // Increase the number of times this state has been executed
+        scenario.Counter++;
+
+        // Only if the number of times this state is executed equals the required StateTimes, 
+        // proceed to next state and reset the counter to 0
+        if (scenario.Counter == (Mapping.TimesInSameState ?? 1))
+        {
+            scenario.NextState = Mapping.NextState;
+            scenario.Counter = 0;
+        }
+
+        // Else just update Started and Finished
+        scenario.Started = true;
+        scenario.Finished = Mapping.NextState == null;
+    }
+
+    internal void LogWebSocketMessage(
         WebSocketMessageDirection direction,
         WebSocketMessageType messageType,
         object? data,
@@ -150,7 +194,7 @@ public class WireMockWebSocketContext : IWebSocketContext
             bodyData = new BodyData
             {
                 BodyAsString = textContent,
-                DetectedBodyType = BodyType.String
+                DetectedBodyType = BodyType.WebSocketText
             };
         }
         else if (messageType == WebSocketMessageType.Binary && data is byte[] binary)
@@ -158,7 +202,7 @@ public class WireMockWebSocketContext : IWebSocketContext
             bodyData = new BodyData
             {
                 BodyAsBytes = binary,
-                DetectedBodyType = BodyType.Bytes
+                DetectedBodyType = BodyType.WebSocketBinary
             };
         }
         else
@@ -166,7 +210,7 @@ public class WireMockWebSocketContext : IWebSocketContext
             bodyData = new BodyData
             {
                 BodyAsString = messageType.ToString(),
-                DetectedBodyType = BodyType.String
+                DetectedBodyType = BodyType.WebSocketClose
             };
         }
 
@@ -225,93 +269,47 @@ public class WireMockWebSocketContext : IWebSocketContext
         activity?.Dispose();
     }
 
-    /// <inheritdoc />
-    public async Task CloseAsync(WebSocketCloseStatus closeStatus, string statusDescription)
+    private async Task SendAsyncInternal(
+        ArraySegment<byte> buffer,
+        WebSocketMessageType messageType,
+        bool endOfMessage,
+        object? data,
+        CancellationToken cancellationToken)
     {
-        await WebSocket.CloseAsync(closeStatus, statusDescription, CancellationToken.None);
+        Activity? activity = null;
+        var shouldTrace = Options.ActivityTracingOptions is not null;
 
-        LogWebSocketMessage(WebSocketMessageDirection.Send, WebSocketMessageType.Close, $"CloseStatus: {closeStatus}, Description: {statusDescription}", null);
-    }
-
-    /// <inheritdoc />
-    public void SetScenarioState(string nextState)
-    {
-        SetScenarioState(nextState, null);
-    }
-
-    /// <inheritdoc />
-    public void SetScenarioState(string nextState, string? description)
-    {
-        if (Mapping.Scenario == null)
+        if (shouldTrace)
         {
-            return;
+            activity = WireMockActivitySource.StartWebSocketMessageActivity(WebSocketMessageDirection.Send, Mapping.Guid);
+            WireMockActivitySource.EnrichWithWebSocketMessage(
+                activity,
+                messageType,
+                buffer.Count,
+                endOfMessage,
+                data as string,
+                Options.ActivityTracingOptions
+            );
         }
 
-        // Use the same logic as WireMockMiddleware
-        if (Options.Scenarios.TryGetValue(Mapping.Scenario, out var scenarioState))
+        try
         {
-            // Directly set the next state (bypass counter logic for manual WebSocket state changes)
-            scenarioState.NextState = nextState;
-            scenarioState.Started = true;
-            scenarioState.Finished = nextState == null;
+            await WebSocket.SendAsync(buffer, messageType, endOfMessage, cancellationToken).ConfigureAwait(false);
 
-            // Reset counter when manually setting state
-            scenarioState.Counter = 0;
-        }
-        else
-        {
-            // Create new scenario state if it doesn't exist
-            Options.Scenarios.TryAdd(Mapping.Scenario, new ScenarioState
+            // Log the send operation
+            if (Options.MaxRequestLogCount is null or > 0)
             {
-                Name = Mapping.Scenario,
-                NextState = nextState,
-                Started = true,
-                Finished = nextState == null,
-                Counter = 0
-            });
+                LogWebSocketMessage(WebSocketMessageDirection.Send, messageType, data, activity);
+            }
         }
-    }
-
-    /// <summary>
-    /// Update scenario state following the same pattern as WireMockMiddleware.UpdateScenarioState
-    /// This is called automatically when the WebSocket connection is established.
-    /// </summary>
-    internal void UpdateScenarioState()
-    {
-        if (Mapping.Scenario == null)
+        catch (Exception ex)
         {
-            return;
+            WireMockActivitySource.RecordException(activity, ex);
+            throw;
         }
-
-        // Ensure scenario exists
-        if (!Options.Scenarios.TryGetValue(Mapping.Scenario, out var scenario))
+        finally
         {
-            return;
-        }
-
-        // Follow exact same logic as WireMockMiddleware.UpdateScenarioState
-        // Increase the number of times this state has been executed
-        scenario.Counter++;
-
-        // Only if the number of times this state is executed equals the required StateTimes, 
-        // proceed to next state and reset the counter to 0
-        if (scenario.Counter == (Mapping.TimesInSameState ?? 1))
-        {
-            scenario.NextState = Mapping.NextState;
-            scenario.Counter = 0;
-        }
-
-        // Else just update Started and Finished
-        scenario.Started = true;
-        scenario.Finished = Mapping.NextState == null;
-    }
-
-    /// <inheritdoc />
-    public async Task BroadcastTextAsync(string text, CancellationToken cancellationToken = default)
-    {
-        if (Registry != null)
-        {
-            await Registry.BroadcastTextAsync(text, cancellationToken);
+            activity?.Dispose();
         }
     }
 }