websocket-relay/internal/hub/hub_edge_cases_test.go

package hub

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

// TestIntegration_PathWithSlashes verifies that clients connected to a nested
// path (e.g., /a/b/c) can communicate within the same room.
func TestIntegration_PathWithSlashes(t *testing.T) {
	server, h := setupTestServer(t)
	defer server.Close()
	defer h.Shutdown()

	// Connect two clients to a nested path
	conn1 := dialWSPath(t, server, "/a/b/c")
	defer conn1.Close()
	conn2 := dialWSPath(t, server, "/a/b/c")
	defer conn2.Close()

	waitForClients(t, h, 2, time.Second)

	// Verify they are in the same room
	testMsg := "nested path message"
	if err := conn1.WriteMessage(websocket.TextMessage, []byte(testMsg)); err != nil {
		t.Fatalf("Failed to send message: %v", err)
	}

	// Both clients should receive it
	for i, conn := range []*websocket.Conn{conn1, conn2} {
		conn.SetReadDeadline(time.Now().Add(time.Second))
		_, msg, err := conn.ReadMessage()
		if err != nil {
			t.Fatalf("Client %d failed to read message: %v", i+1, err)
		}
		if string(msg) != testMsg {
			t.Errorf("Client %d expected %q, got %q", i+1, testMsg, string(msg))
		}
	}

	// Verify a different nested path is a separate room
	conn3 := dialWSPath(t, server, "/a/b/d")
	defer conn3.Close()

	waitForClients(t, h, 3, time.Second)

	// Send from /a/b/d
	otherMsg := "different path message"
	if err := conn3.WriteMessage(websocket.TextMessage, []byte(otherMsg)); err != nil {
		t.Fatalf("Failed to send message from /a/b/d: %v", err)
	}

	// conn3 should get its own message
	conn3.SetReadDeadline(time.Now().Add(time.Second))
	_, msg, err := conn3.ReadMessage()
	if err != nil {
		t.Fatalf("Client 3 failed to read own message: %v", err)
	}
	if string(msg) != otherMsg {
		t.Errorf("Client 3 expected %q, got %q", otherMsg, string(msg))
	}

	// conn1 and conn2 should NOT receive it
	for i, conn := range []*websocket.Conn{conn1, conn2} {
		conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
		_, _, err := conn.ReadMessage()
		if err == nil {
			t.Errorf("Client %d in /a/b/c should NOT have received message from /a/b/d", i+1)
		}
	}
}

// TestIntegration_QueryStringIgnored verifies that clients connecting to the
// same path with different query strings are placed in the same room.
func TestIntegration_QueryStringIgnored(t *testing.T) {
	server, h := setupTestServer(t)
	defer server.Close()
	defer h.Shutdown()

	// Connect client 1 to /room?token=abc
	conn1 := dialWSPath(t, server, "/room?token=abc")
	defer conn1.Close()

	// Connect client 2 to /room?token=xyz
	conn2 := dialWSPath(t, server, "/room?token=xyz")
	defer conn2.Close()

	waitForClients(t, h, 2, time.Second)

	// They should be in the same room (/room) since query strings are stripped
	testMsg := "query string test"
	if err := conn1.WriteMessage(websocket.TextMessage, []byte(testMsg)); err != nil {
		t.Fatalf("Failed to send message: %v", err)
	}

	// Both clients should receive the message
	for i, conn := range []*websocket.Conn{conn1, conn2} {
		conn.SetReadDeadline(time.Now().Add(time.Second))
		_, msg, err := conn.ReadMessage()
		if err != nil {
			t.Fatalf("Client %d failed to read message: %v", i+1, err)
		}
		if string(msg) != testMsg {
			t.Errorf("Client %d expected %q, got %q", i+1, testMsg, string(msg))
		}
	}

	// Verify they are actually in the same room (only 1 room exists)
	if count := h.RoomCount(); count != 1 {
		t.Errorf("Expected 1 room (query strings ignored), got %d", count)
	}
}

// TestIntegration_DefaultRoom verifies that clients connecting to the bare
// root path (/) are placed in the default room and can communicate.
func TestIntegration_DefaultRoom(t *testing.T) {
	server, h := setupTestServer(t)
	defer server.Close()
	defer h.Shutdown()

	// Connect clients to "/" (bare root)
	conn1 := dialWSPath(t, server, "/")
	defer conn1.Close()
	conn2 := dialWSPath(t, server, "/")
	defer conn2.Close()

	waitForClients(t, h, 2, time.Second)

	// Verify broadcast works
	testMsg := "default room message"
	if err := conn1.WriteMessage(websocket.TextMessage, []byte(testMsg)); err != nil {
		t.Fatalf("Failed to send message: %v", err)
	}

	for i, conn := range []*websocket.Conn{conn1, conn2} {
		conn.SetReadDeadline(time.Now().Add(time.Second))
		_, msg, err := conn.ReadMessage()
		if err != nil {
			t.Fatalf("Client %d failed to read message: %v", i+1, err)
		}
		if string(msg) != testMsg {
			t.Errorf("Client %d expected %q, got %q", i+1, testMsg, string(msg))
		}
	}

	// Verify they are in the same room
	if count := h.RoomCount(); count != 1 {
		t.Errorf("Expected 1 room for default root path, got %d", count)
	}
}

// TestIntegration_ClientDisconnectFromRoom verifies that when one client
// disconnects from a room, the remaining clients can still communicate.
func TestIntegration_ClientDisconnectFromRoom(t *testing.T) {
	server, h := setupTestServer(t)
	defer server.Close()
	defer h.Shutdown()

	// Connect 3 clients to the same room
	conn1 := dialWSPath(t, server, "/game")
	defer conn1.Close()
	conn2 := dialWSPath(t, server, "/game")
	defer conn2.Close()
	conn3 := dialWSPath(t, server, "/game")
	defer conn3.Close()

	waitForClients(t, h, 3, time.Second)

	// Disconnect client 1
	conn1.Close()
	waitForClients(t, h, 2, time.Second)

	// Remaining 2 clients should still communicate
	testMsg := "still alive"
	if err := conn2.WriteMessage(websocket.TextMessage, []byte(testMsg)); err != nil {
		t.Fatalf("Failed to send message from client 2: %v", err)
	}

	// conn2 and conn3 should receive the message
	for i, conn := range []*websocket.Conn{conn2, conn3} {
		conn.SetReadDeadline(time.Now().Add(time.Second))
		_, msg, err := conn.ReadMessage()
		if err != nil {
			t.Fatalf("Client %d failed to read message after disconnect: %v", i+2, err)
		}
		if string(msg) != testMsg {
			t.Errorf("Client %d expected %q, got %q", i+2, testMsg, string(msg))
		}
	}

	// Verify the room still exists
	if count := h.RoomCount(); count != 1 {
		t.Errorf("Expected room to still exist with 2 clients, got %d rooms", count)
	}

	// Send from client 3 to confirm bidirectional communication
	replyMsg := "reply from client 3"
	if err := conn3.WriteMessage(websocket.TextMessage, []byte(replyMsg)); err != nil {
		t.Fatalf("Failed to send reply from client 3: %v", err)
	}

	for i, conn := range []*websocket.Conn{conn2, conn3} {
		conn.SetReadDeadline(time.Now().Add(time.Second))
		_, msg, err := conn.ReadMessage()
		if err != nil {
			t.Fatalf("Client %d failed to read reply: %v", i+2, err)
		}
		if string(msg) != replyMsg {
			t.Errorf("Client %d expected %q, got %q", i+2, replyMsg, string(msg))
		}
	}
}

// TestIntegration_ConcurrentRoomOperations verifies that rapidly connecting and
// disconnecting clients across multiple rooms concurrently causes no data races.
func TestIntegration_ConcurrentRoomOperations(t *testing.T) {
	server, h := setupTestServer(t)
	defer server.Close()
	defer h.Shutdown()

	const numRooms = 10
	const clientsPerRoom = 5

	var wg sync.WaitGroup

	// Rapidly connect and disconnect clients across many rooms concurrently
	for i := 0; i < numRooms; i++ {
		wg.Add(1)
		go func(roomIdx int) {
			defer wg.Done()
			room := fmt.Sprintf("/concurrent-room-%d", roomIdx)

			conns := make([]*websocket.Conn, 0, clientsPerRoom)

			// Connect all clients in this room
			for j := 0; j < clientsPerRoom; j++ {
				conn := dialWSPath(t, server, room)
				conns = append(conns, conn)
			}

			// Small delay to allow registration to process
			time.Sleep(50 * time.Millisecond)

			// Send a message from the first client
			if len(conns) > 0 {
				msg := fmt.Sprintf("msg from room %d", roomIdx)
				conns[0].WriteMessage(websocket.TextMessage, []byte(msg))
			}

			// Small delay then disconnect all
			time.Sleep(50 * time.Millisecond)

			for _, conn := range conns {
				conn.Close()
			}
		}(i)
	}

	wg.Wait()

	// Wait for all cleanups to complete
	waitForClients(t, h, 0, 5*time.Second)
	waitForRooms(t, h, 0, 5*time.Second)

	// Verify clean state: no rooms, no clients
	if count := h.ClientCount(); count != 0 {
		t.Errorf("Expected 0 clients after concurrent operations, got %d", count)
	}
	if count := h.RoomCount(); count != 0 {
		t.Errorf("Expected 0 rooms after concurrent operations, got %d", count)
	}
}