44"context"
55"net/http"
66"net/http/httptest"
7- "sync"
87"testing"
98"time"
109
@@ -96,16 +95,17 @@ func TestLoggerMiddleware_WebSocket(t *testing.T) {
9695sink := & fakeSink {}
9796logger := slog .Make (sink )
9897logger = logger .Leveled (slog .LevelDebug )
99- var wg sync. WaitGroup
98+
10099// Create a test handler to simulate a WebSocket connection
101100testHandler := http .HandlerFunc (func (rw http.ResponseWriter ,r * http.Request ) {
102101conn ,err := websocket .Accept (rw ,r ,nil )
103102if err != nil {
104103t .Errorf ("failed to accept websocket: %v" ,err )
105104return
106105}
106+ requestLgr := RequestLoggerFromContext (r .Context ())
107+ requestLgr .WriteLog (r .Context (),http .StatusSwitchingProtocols )
107108defer conn .Close (websocket .StatusNormalClosure ,"" )
108- defer wg .Done ()
109109
110110// Send a couple of messages for testing
111111_ = conn .Write (ctx ,websocket .MessageText , []byte ("ping" ))
@@ -125,10 +125,8 @@ func TestLoggerMiddleware_WebSocket(t *testing.T) {
125125// Create a test HTTP request
126126srv := httptest .NewServer (customHandler )
127127defer srv .Close ()
128- wg .Add (1 )
129128// nolint: bodyclose
130129conn ,_ ,err := websocket .Dial (ctx ,srv .URL ,nil )
131- wg .Wait ()
132130if err != nil {
133131t .Fatalf ("failed to create WebSocket connection: %v" ,err )
134132}