src/server/internal/ws/hub_test.go

2,860 bytes · 142 lines · capsule://quake0day/[email protected] raw on github

package ws

import (
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

func newTestClient(hub *Hub, sessionID string) *Client {
	return &Client{
		SessionID: sessionID,
		Conn:      &websocket.Conn{}, // nil-safe for non-network tests
		Send:      make(chan []byte, 64),
		hub:       hub,
	}
}

func TestHubRegisterUnregister(t *testing.T) {
	hub := NewHub()
	client := newTestClient(hub, "session-1")

	hub.Register(client)

	hub.mu.RLock()
	clients, ok := hub.clients["session-1"]
	hub.mu.RUnlock()

	if !ok || len(clients) != 1 {
		t.Fatal("expected 1 client registered for session-1")
	}

	hub.Unregister(client)

	hub.mu.RLock()
	_, ok = hub.clients["session-1"]
	hub.mu.RUnlock()

	if ok {
		t.Error("expected session-1 to be removed after last client unregisters")
	}
}

func TestHubBroadcast(t *testing.T) {
	hub := NewHub()
	c1 := newTestClient(hub, "s1")
	c2 := newTestClient(hub, "s1")
	c3 := newTestClient(hub, "s2") // different session

	hub.Register(c1)
	hub.Register(c2)
	hub.Register(c3)

	msg := []byte(`{"type":"test"}`)
	hub.Broadcast("s1", msg)

	// c1 and c2 should receive the message
	select {
	case received := <-c1.Send:
		if string(received) != string(msg) {
			t.Errorf("c1 got %s, want %s", received, msg)
		}
	case <-time.After(time.Second):
		t.Error("c1 did not receive message")
	}

	select {
	case received := <-c2.Send:
		if string(received) != string(msg) {
			t.Errorf("c2 got %s, want %s", received, msg)
		}
	case <-time.After(time.Second):
		t.Error("c2 did not receive message")
	}

	// c3 should not receive the message
	select {
	case <-c3.Send:
		t.Error("c3 should not receive message from different session")
	case <-time.After(50 * time.Millisecond):
		// expected
	}

	// Cleanup
	hub.Unregister(c1)
	hub.Unregister(c2)
	hub.Unregister(c3)
}

func TestHubBroadcastNoSession(t *testing.T) {
	hub := NewHub()
	// Should not panic
	hub.Broadcast("nonexistent", []byte("hello"))
}

func TestHubBroadcastJSON(t *testing.T) {
	hub := NewHub()
	c := newTestClient(hub, "s1")
	hub.Register(c)

	hub.BroadcastJSON("s1", map[string]string{"type": "test", "data": "hello"})

	select {
	case msg := <-c.Send:
		if len(msg) == 0 {
			t.Error("expected non-empty JSON message")
		}
	case <-time.After(time.Second):
		t.Error("did not receive JSON broadcast")
	}

	hub.Unregister(c)
}

func TestHubMultipleSessionsIsolated(t *testing.T) {
	hub := NewHub()
	c1 := newTestClient(hub, "session-a")
	c2 := newTestClient(hub, "session-b")

	hub.Register(c1)
	hub.Register(c2)

	hub.Broadcast("session-a", []byte("for-a"))

	select {
	case <-c1.Send:
		// expected
	case <-time.After(time.Second):
		t.Error("c1 should receive message")
	}

	select {
	case <-c2.Send:
		t.Error("c2 should not receive message for session-a")
	case <-time.After(50 * time.Millisecond):
		// expected
	}

	hub.Unregister(c1)
	hub.Unregister(c2)
}