package api import ( "context" "encoding/json" "net/http/httptest" "strings" "testing" "time" "github.com/go-chi/chi/v5" "github.com/gorilla/websocket" "git.threesix.ai/jordan/sp3-test-1770368381/pkg/logging" "git.threesix.ai/jordan/sp3-test-1770368381/pkg/realtime" ) func TestWebSocketUpgrade(t *testing.T) { logger := logging.Nop() hub := realtime.NewHub(logger) ctx, cancel := context.WithCancel(context.Background()) defer cancel() go hub.Run(ctx) wsHandler := realtime.NewHandler(hub, logger, realtime.HandlerConfig{}) r := chi.NewRouter() r.Mount("/api/chat-api/ws", wsHandler.Routes()) server := httptest.NewServer(r) defer server.Close() wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/api/chat-api/ws" conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) if err != nil { t.Fatalf("failed to connect to websocket: %v", err) } defer conn.Close() if resp.StatusCode != 101 { t.Errorf("expected status 101, got %d", resp.StatusCode) } } func TestWebSocketBroadcast(t *testing.T) { logger := logging.Nop() hub := realtime.NewHub(logger) ctx, cancel := context.WithCancel(context.Background()) defer cancel() go hub.Run(ctx) wsHandler := realtime.NewHandler(hub, logger, realtime.HandlerConfig{}) r := chi.NewRouter() r.Mount("/api/chat-api/ws", wsHandler.Routes()) server := httptest.NewServer(r) defer server.Close() wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/api/chat-api/ws" // Connect two clients conn1, _, err := websocket.DefaultDialer.Dial(wsURL, nil) if err != nil { t.Fatalf("client 1 failed to connect: %v", err) } defer conn1.Close() conn2, _, err := websocket.DefaultDialer.Dial(wsURL, nil) if err != nil { t.Fatalf("client 2 failed to connect: %v", err) } defer conn2.Close() // Wait for both clients to register with the hub time.Sleep(100 * time.Millisecond) // Client 1 sends a message msg := realtime.Message{ Type: "chat", Data: json.RawMessage(`"hello from client 1"`), } if err := conn1.WriteJSON(msg); err != nil { t.Fatalf("client 1 failed to send: %v", err) } // Client 2 should receive the broadcast conn2.SetReadDeadline(time.Now().Add(2 * time.Second)) var received realtime.Message if err := conn2.ReadJSON(&received); err != nil { t.Fatalf("client 2 failed to receive: %v", err) } if received.Type != "chat" { t.Errorf("expected message type 'chat', got %q", received.Type) } var data string if err := json.Unmarshal(received.Data, &data); err != nil { t.Fatalf("failed to unmarshal data: %v", err) } if data != "hello from client 1" { t.Errorf("expected 'hello from client 1', got %q", data) } } func TestWebSocketRoomBroadcast(t *testing.T) { logger := logging.Nop() hub := realtime.NewHub(logger) ctx, cancel := context.WithCancel(context.Background()) defer cancel() go hub.Run(ctx) wsHandler := realtime.NewHandler(hub, logger, realtime.HandlerConfig{}) r := chi.NewRouter() r.Mount("/api/chat-api/ws", wsHandler.Routes()) server := httptest.NewServer(r) defer server.Close() baseURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/api/chat-api/ws" // Connect client A to room "general" connA, _, err := websocket.DefaultDialer.Dial(baseURL+"/general", nil) if err != nil { t.Fatalf("client A failed to connect: %v", err) } defer connA.Close() // Connect client B to room "general" connB, _, err := websocket.DefaultDialer.Dial(baseURL+"/general", nil) if err != nil { t.Fatalf("client B failed to connect: %v", err) } defer connB.Close() // Connect client C to room "other" connC, _, err := websocket.DefaultDialer.Dial(baseURL+"/other", nil) if err != nil { t.Fatalf("client C failed to connect: %v", err) } defer connC.Close() // Wait for all clients to register time.Sleep(100 * time.Millisecond) // Client A sends a message (room defaults to "general") msg := realtime.Message{ Type: "chat", Data: json.RawMessage(`"room message"`), } if err := connA.WriteJSON(msg); err != nil { t.Fatalf("client A failed to send: %v", err) } // Client B should receive it (same room) connB.SetReadDeadline(time.Now().Add(2 * time.Second)) var received realtime.Message if err := connB.ReadJSON(&received); err != nil { t.Fatalf("client B failed to receive: %v", err) } if received.Type != "chat" { t.Errorf("expected type 'chat', got %q", received.Type) } // Client C should NOT receive it (different room) connC.SetReadDeadline(time.Now().Add(200 * time.Millisecond)) var unexpected realtime.Message err = connC.ReadJSON(&unexpected) if err == nil { t.Error("client C should not have received message from 'general' room") } }