44"context"
55"net/http"
66"net/http/httptest"
7+ "sync"
78"testing"
89"time"
910
@@ -95,7 +96,7 @@ func TestLoggerMiddleware_WebSocket(t *testing.T) {
9596sink := & fakeSink {}
9697logger := slog .Make (sink )
9798logger = logger .Leveled (slog .LevelDebug )
98-
99+ wg := sync. WaitGroup {}
99100// Create a test handler to simulate a WebSocket connection
100101testHandler := http .HandlerFunc (func (rw http.ResponseWriter ,r * http.Request ) {
101102conn ,err := websocket .Accept (rw ,r ,nil )
@@ -105,6 +106,7 @@ func TestLoggerMiddleware_WebSocket(t *testing.T) {
105106}
106107requestLgr := RequestLoggerFromContext (r .Context ())
107108requestLgr .WriteLog (r .Context (),http .StatusSwitchingProtocols )
109+ wg .Done ()
108110defer conn .Close (websocket .StatusNormalClosure ,"" )
109111
110112// Send a couple of messages for testing
@@ -122,16 +124,16 @@ func TestLoggerMiddleware_WebSocket(t *testing.T) {
122124wrappedHandler .ServeHTTP (sw ,r )
123125})
124126
125- // Create a test HTTP request
126127srv := httptest .NewServer (customHandler )
127128defer srv .Close ()
129+ wg .Add (1 )
128130// nolint: bodyclose
129131conn ,_ ,err := websocket .Dial (ctx ,srv .URL ,nil )
130132if err != nil {
131133t .Fatalf ("failed to create WebSocket connection: %v" ,err )
132134}
133135defer conn .Close (websocket .StatusNormalClosure ,"" )
134-
136+ wg . Wait ()
135137if len (sink .entries )!= 1 {
136138t .Fatalf ("expected 1 log entry, got %d" ,len (sink .entries ))
137139}