added WebSocket test to verify early logging

ibetitsmike · ibetitsmike · commit dfa2f4fc2f1c · 2025-04-07T09:45:47.000Z
diff --git a/coderd/httpmw/logger_internal_test.go b/coderd/httpmw/logger_internal_test.go
@@ -8,7 +8,10 @@ import (
 	"time"
 
 	"cdr.dev/slog"
+	"github.com/coder/coder/v2/coderd/httpapi"
 	"github.com/coder/coder/v2/coderd/tracing"
+	"github.com/coder/coder/v2/testutil"
+	"github.com/coder/websocket"
 )
 
 func TestRequestLogger_WriteLog(t *testing.T) {
@@ -48,7 +51,7 @@ func TestRequestLogger_WriteLog(t *testing.T) {
 	}
 }
 
-func TestLoggerMiddleware(t *testing.T) {
+func TestLoggerMiddleware_SingleRequest(t *testing.T) {
 	t.Parallel()
 
 	sink := &fakeSink{}
@@ -85,6 +88,54 @@ func TestLoggerMiddleware(t *testing.T) {
 	}
 }
 
+func TestLoggerMiddleware_WebSocket(t *testing.T) {
+	t.Parallel()
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
+
+	sink := &fakeSink{}
+	logger := slog.Make(sink)
+	logger = logger.Leveled(slog.LevelDebug)
+
+	// Create a test handler to simulate a WebSocket connection
+	testHandler := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		_, err := websocket.Accept(rw, r, nil)
+		if err != nil {
+			httpapi.Write(ctx, rw, http.StatusBadRequest, nil)
+			return
+		}
+		time.Sleep(1000)
+	})
+
+	// Wrap the test handler with the Logger middleware
+	loggerMiddleware := Logger(logger)
+	wrappedHandler := loggerMiddleware(testHandler)
+
+	// RequestLogger expects the ResponseWriter to be *tracing.StatusWriter
+	customHandler := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		sw := &tracing.StatusWriter{ResponseWriter: rw}
+		wrappedHandler.ServeHTTP(sw, r)
+	})
+
+	// Create a test HTTP request
+	srv := httptest.NewServer(customHandler)
+	defer srv.Close()
+
+	conn, _, err := websocket.Dial(ctx, srv.URL, nil)
+	if err != nil {
+		t.Fatalf("failed to create WebSocket connection: %v", err)
+	}
+	defer conn.Close(websocket.StatusNormalClosure, "")
+
+	if len(sink.entries) != 1 {
+		t.Fatalf("expected 1 log entry, got %d", len(sink.entries))
+	}
+
+	if sink.entries[0].Message != "GET" {
+		t.Errorf("expected log message to be 'GET', got '%s'", sink.entries[0].Message)
+	}
+}
+
 type fakeSink struct {
 	entries []slog.SinkEntry
 }
