test(hub): add integration tests for room isolation
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
This commit is contained in:
savinmax
parent
c226bdab7f
commit
516f8c5008
293
internal/hub/hub_room_isolation_test.go
Normal file
293
internal/hub/hub_room_isolation_test.go
Normal file
@ -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)
|
||||
}
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user