From 516f8c50089232f87319929d27aec05de59ccdd7 Mon Sep 17 00:00:00 2001 From: savinmax Date: Sat, 13 Jun 2026 13:30:43 +0200 Subject: [PATCH] test(hub): add integration tests for room isolation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add dedicated integration tests verifying message isolation between rooms: - TestIntegration_SameRoomBroadcast: two clients in /chat both receive broadcast messages - TestIntegration_CrossRoomIsolation: client in /room-a does not leak messages to client in /room-b (verified via read deadline timeout) - TestIntegration_MultipleRoomsSimultaneous: 3 rooms with 2 clients each, messages stay within their room - TestIntegration_RoomCleanup: verifies RoomCount() increases on connect and decreases (room removed) on last client disconnect - TestIntegration_RoomCleanup_MultipleClients: room persists while any client remains connected - TestIntegration_RoomCleanup_ConcurrentDisconnects: 5 rooms cleaned up concurrently without races Also introduces dialWSPath(t, server, path) helper for tests that specify WebSocket paths directly with a leading slash. All tests pass with -race flag. 🤖 Assisted by the code-assist SOP --- internal/hub/hub_room_isolation_test.go | 293 ++++++++++++++++++++++++ 1 file changed, 293 insertions(+) create mode 100644 internal/hub/hub_room_isolation_test.go diff --git a/internal/hub/hub_room_isolation_test.go b/internal/hub/hub_room_isolation_test.go new file mode 100644 index 0000000..b1426cd --- /dev/null +++ b/internal/hub/hub_room_isolation_test.go @@ -0,0 +1,293 @@ +package hub + +import ( + "bytes" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/gorilla/websocket" + "websocket-relay/internal/logging" +) + +// dialWSPath dials a WebSocket connection to the test server at the given path. +// The path should include a leading slash (e.g., "/chat", "/room-a"). +func dialWSPath(t *testing.T, server *httptest.Server, path string) *websocket.Conn { + t.Helper() + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + path + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("Failed to dial WebSocket (path=%q): %v", path, err) + } + return conn +} + +// helper: wait until hub reaches expected room count or timeout +func waitForRooms(t *testing.T, h *Hub, expected int, timeout time.Duration) { + t.Helper() + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if h.RoomCount() == expected { + return + } + time.Sleep(5 * time.Millisecond) + } + t.Fatalf("Timed out waiting for %d rooms, got %d", expected, h.RoomCount()) +} + +func TestIntegration_SameRoomBroadcast(t *testing.T) { + server, h := setupTestServer(t) + defer server.Close() + defer h.Shutdown() + + // Connect 2 clients to /chat + conn1 := dialWSPath(t, server, "/chat") + defer conn1.Close() + conn2 := dialWSPath(t, server, "/chat") + defer conn2.Close() + + waitForClients(t, h, 2, time.Second) + + // Client 1 sends message + testMsg := "hello chat room" + 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)) + } + } +} + +func TestIntegration_CrossRoomIsolation(t *testing.T) { + server, h := setupTestServer(t) + defer server.Close() + defer h.Shutdown() + + // Connect client A to /room-a, client B to /room-b + connA := dialWSPath(t, server, "/room-a") + defer connA.Close() + connB := dialWSPath(t, server, "/room-b") + defer connB.Close() + + waitForClients(t, h, 2, time.Second) + + // Client A sends message + testMsg := "message from room-a" + if err := connA.WriteMessage(websocket.TextMessage, []byte(testMsg)); err != nil { + t.Fatalf("Failed to send message from client A: %v", err) + } + + // Client A receives it (echo to self within room) + connA.SetReadDeadline(time.Now().Add(time.Second)) + _, msg, err := connA.ReadMessage() + if err != nil { + t.Fatalf("Client A failed to read own message: %v", err) + } + if string(msg) != testMsg { + t.Errorf("Client A expected %q, got %q", testMsg, string(msg)) + } + + // Client B does NOT receive it (verify with read deadline timeout) + connB.SetReadDeadline(time.Now().Add(100 * time.Millisecond)) + _, _, err = connB.ReadMessage() + if err == nil { + t.Fatal("Client B should NOT have received a message from room-a") + } + // Timeout error is expected — message was correctly isolated +} + +func TestIntegration_MultipleRoomsSimultaneous(t *testing.T) { + server, h := setupTestServer(t) + defer server.Close() + defer h.Shutdown() + + // 3 rooms with 2 clients each + type roomClients struct { + path string + conns []*websocket.Conn + } + + rooms := []roomClients{ + {path: "/room-1"}, + {path: "/room-2"}, + {path: "/room-3"}, + } + + for i := range rooms { + for j := 0; j < 2; j++ { + conn := dialWSPath(t, server, rooms[i].path) + defer conn.Close() + rooms[i].conns = append(rooms[i].conns, conn) + } + } + + waitForClients(t, h, 6, time.Second) + waitForRooms(t, h, 3, time.Second) + + // Send a message in each room from the first client + for i, room := range rooms { + msg := fmt.Sprintf("message for %s", room.path) + if err := room.conns[0].WriteMessage(websocket.TextMessage, []byte(msg)); err != nil { + t.Fatalf("Room %d: failed to send message: %v", i+1, err) + } + } + + // Verify each room's clients receive only their room's message + for i, room := range rooms { + expectedMsg := fmt.Sprintf("message for %s", room.path) + for j, conn := range room.conns { + conn.SetReadDeadline(time.Now().Add(time.Second)) + _, msg, err := conn.ReadMessage() + if err != nil { + t.Fatalf("Room %d, client %d: failed to read message: %v", i+1, j+1, err) + } + if string(msg) != expectedMsg { + t.Errorf("Room %d, client %d: expected %q, got %q", i+1, j+1, expectedMsg, string(msg)) + } + } + } + + // Verify cross-room isolation: send another message from room-1 and confirm + // room-2 and room-3 clients don't receive it + isolationMsg := "room-1 only" + if err := rooms[0].conns[0].WriteMessage(websocket.TextMessage, []byte(isolationMsg)); err != nil { + t.Fatalf("Failed to send isolation test message: %v", err) + } + + // Room-1 clients should get it + for j, conn := range rooms[0].conns { + conn.SetReadDeadline(time.Now().Add(time.Second)) + _, msg, err := conn.ReadMessage() + if err != nil { + t.Fatalf("Room-1, client %d: failed to read isolation message: %v", j+1, err) + } + if string(msg) != isolationMsg { + t.Errorf("Room-1, client %d: expected %q, got %q", j+1, isolationMsg, string(msg)) + } + } + + // Room-2 and room-3 clients should NOT get it + for i := 1; i < len(rooms); i++ { + for j, conn := range rooms[i].conns { + conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond)) + _, _, err := conn.ReadMessage() + if err == nil { + t.Errorf("Room %d, client %d: should NOT have received room-1's message", i+1, j+1) + } + } + } +} + +func TestIntegration_RoomCleanup(t *testing.T) { + logger := logging.NewLogger("debug", &bytes.Buffer{}) + h := New(logger) + go h.Run() + defer h.Shutdown() + + server := httptest.NewServer(http.HandlerFunc(h.HandleWebSocket)) + defer server.Close() + + // Record initial room count + initialRooms := h.RoomCount() + + // Connect client to /temp + conn := dialWSPath(t, server, "/temp") + + waitForClients(t, h, 1, time.Second) + waitForRooms(t, h, initialRooms+1, time.Second) + + // Verify RoomCount increased + if count := h.RoomCount(); count != initialRooms+1 { + t.Errorf("Expected room count to be %d, got %d", initialRooms+1, count) + } + + // Disconnect the client + conn.Close() + + // Wait for cleanup + waitForClients(t, h, 0, time.Second) + waitForRooms(t, h, initialRooms, time.Second) + + // Verify RoomCount decreased (room removed) + if count := h.RoomCount(); count != initialRooms { + t.Errorf("Expected room count to return to %d after disconnect, got %d", initialRooms, count) + } +} + +func TestIntegration_RoomCleanup_MultipleClients(t *testing.T) { + server, h := setupTestServer(t) + defer server.Close() + defer h.Shutdown() + + // Connect two clients to the same room + conn1 := dialWSPath(t, server, "/shared") + conn2 := dialWSPath(t, server, "/shared") + + waitForClients(t, h, 2, time.Second) + waitForRooms(t, h, 1, time.Second) + + // Disconnect first client — room should still exist + conn1.Close() + waitForClients(t, h, 1, time.Second) + + if count := h.RoomCount(); count != 1 { + t.Errorf("Expected room to still exist with 1 client remaining, got %d rooms", count) + } + + // Disconnect second client — room should be cleaned up + conn2.Close() + waitForClients(t, h, 0, time.Second) + waitForRooms(t, h, 0, time.Second) + + if count := h.RoomCount(); count != 0 { + t.Errorf("Expected room to be removed after all clients disconnect, got %d rooms", count) + } +} + +func TestIntegration_RoomCleanup_ConcurrentDisconnects(t *testing.T) { + server, h := setupTestServer(t) + defer server.Close() + defer h.Shutdown() + + const numRooms = 5 + conns := make([]*websocket.Conn, numRooms) + for i := 0; i < numRooms; i++ { + path := fmt.Sprintf("/room-%d", i) + conns[i] = dialWSPath(t, server, path) + } + + waitForClients(t, h, numRooms, time.Second) + waitForRooms(t, h, numRooms, time.Second) + + // Disconnect all clients concurrently + var wg sync.WaitGroup + for _, conn := range conns { + wg.Add(1) + go func(c *websocket.Conn) { + defer wg.Done() + c.Close() + }(conn) + } + wg.Wait() + + // Wait for all cleanup to complete + waitForClients(t, h, 0, 2*time.Second) + waitForRooms(t, h, 0, 2*time.Second) + + if count := h.RoomCount(); count != 0 { + t.Errorf("Expected 0 rooms after all disconnects, got %d", count) + } +}