sp3-test-1770368381/services/chat-api/internal/api/ws_test.go

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")
	}
}