@@ -4,11 +4,11 @@ import (
44"context"
55"net/http"
66"net/http/httptest"
7+ "sync"
78"testing"
89"time"
910
1011"cdr.dev/slog"
11- "github.com/coder/coder/v2/coderd/httpapi"
1212"github.com/coder/coder/v2/coderd/tracing"
1313"github.com/coder/coder/v2/testutil"
1414"github.com/coder/websocket"
@@ -96,15 +96,22 @@ func TestLoggerMiddleware_WebSocket(t *testing.T) {
9696sink := & fakeSink {}
9797logger := slog .Make (sink )
9898logger = logger .Leveled (slog .LevelDebug )
99-
99+ var wg sync. WaitGroup
100100// Create a test handler to simulate a WebSocket connection
101101testHandler := http .HandlerFunc (func (rw http.ResponseWriter ,r * http.Request ) {
102- _ ,err := websocket .Accept (rw ,r ,nil )
102+
103+ conn ,err := websocket .Accept (rw ,r ,nil )
103104if err != nil {
104- httpapi . Write ( ctx , rw , http . StatusBadRequest , nil )
105+ t . Errorf ( "failed to accept websocket: %v" , err )
105106return
106107}
107- time .Sleep (1000 )
108+ defer conn .Close (websocket .StatusNormalClosure ,"" )
109+ defer wg .Done ()
110+
111+ // Send a couple of messages for testing
112+ _ = conn .Write (ctx ,websocket .MessageText , []byte ("ping" ))
113+ _ = conn .Write (ctx ,websocket .MessageText , []byte ("pong" ))
114+
108115})
109116
110117// Wrap the test handler with the Logger middleware
@@ -120,9 +127,10 @@ func TestLoggerMiddleware_WebSocket(t *testing.T) {
120127// Create a test HTTP request
121128srv := httptest .NewServer (customHandler )
122129defer srv .Close ()
123-
130+ wg . Add ( 1 )
124131// nolint: bodyclose
125132conn ,_ ,err := websocket .Dial (ctx ,srv .URL ,nil )
133+ wg .Wait ()
126134if err != nil {
127135t .Fatalf ("failed to create WebSocket connection: %v" ,err )
128136}