package hub import ( "net/http" "net/http/httptest" "strings" "sync" "testing" "github.com/gorilla/websocket" ) // newTestConn returns a real *websocket.Conn (client side) backed by a local // httptest server. The server-side handler blocks on a done channel until the // test ends; cleanup closes the client conn first, then the server, avoiding // the httptest.Server.Close() hang that occurs when the handler is blocked on // ReadMessage. func newTestConn(t *testing.T) *websocket.Conn { t.Helper() done := make(chan struct{}) srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { up := websocket.Upgrader{} c, err := up.Upgrade(w, r, nil) if err != nil { t.Errorf("upgrade: %v", err) return } <-done _ = c.Close() })) cli, _, err := websocket.DefaultDialer.Dial(strings.Replace(srv.URL, "http:", "ws:", 1), nil) if err != nil { close(done) srv.Close() t.Fatalf("dial: %v", err) } t.Cleanup(func() { close(done) _ = cli.Close() srv.Close() }) return cli } // TestRegisterAndUnregister verifies basic Hub register/unregister bookkeeping. func TestRegisterAndUnregister(t *testing.T) { h := NewHub(5) conn := newTestConn(t) c, err := h.Register("user-1", conn) if err != nil { t.Fatalf("Register: %v", err) } if c.UserID != "user-1" { t.Errorf("UserID = %q, want user-1", c.UserID) } if c.ConnID == "" { t.Error("ConnID should be non-empty") } if h.ActiveConnections() != 1 { t.Errorf("ActiveConnections = %d, want 1", h.ActiveConnections()) } if h.UserCount() != 1 { t.Errorf("UserCount = %d, want 1", h.UserCount()) } if !h.HasUser("user-1") { t.Error("HasUser(user-1) = false, want true") } h.Unregister(c) if h.ActiveConnections() != 0 { t.Errorf("after Unregister ActiveConnections = %d, want 0", h.ActiveConnections()) } if h.HasUser("user-1") { t.Error("HasUser(user-1) = true, want false") } } // TestRegisterTooManyConnections verifies the per-user cap is enforced. func TestRegisterTooManyConnections(t *testing.T) { h := NewHub(2) for i := 0; i < 2; i++ { conn := newTestConn(t) if _, err := h.Register("user-cap", conn); err != nil { t.Fatalf("Register #%d: %v", i, err) } } conn := newTestConn(t) _, err := h.Register("user-cap", conn) if err != ErrTooManyConnections { t.Errorf("Register over cap err = %v, want ErrTooManyConnections", err) } } // TestRegisterWhenClosing verifies that CloseAll rejects further Register. func TestRegisterWhenClosing(t *testing.T) { h := NewHub(5) h.CloseAll() conn := newTestConn(t) _, err := h.Register("user-x", conn) if err != ErrHubClosing { t.Errorf("Register after CloseAll err = %v, want ErrHubClosing", err) } if !h.IsClosing() { t.Error("IsClosing = false, want true") } } // TestSendToUser verifies message delivery to a single user. func TestSendToUser(t *testing.T) { h := NewHub(5) conn := newTestConn(t) c, err := h.Register("user-send", conn) if err != nil { t.Fatalf("Register: %v", err) } defer h.Unregister(c) msg := []byte(`{"type":"message","event":"test"}`) delivered := h.SendToUser("user-send", msg) if delivered != 1 { t.Errorf("SendToUser delivered = %d, want 1", delivered) } // Unknown user: 0 delivered. delivered = h.SendToUser("user-absent", msg) if delivered != 0 { t.Errorf("SendToUser absent delivered = %d, want 0", delivered) } } // TestBroadcast verifies message delivery to all connections. func TestBroadcast(t *testing.T) { h := NewHub(5) conn1 := newTestConn(t) conn2 := newTestConn(t) c1, _ := h.Register("u1", conn1) c2, _ := h.Register("u2", conn2) defer h.Unregister(c1) defer h.Unregister(c2) msg := []byte(`{"type":"message","event":"broadcast"}`) reached := h.Broadcast(msg) if reached != 2 { t.Errorf("Broadcast reached = %d, want 2", reached) } } // TestPresenceHooks verifies that Register/Unregister call the hooks. func TestPresenceHooks(t *testing.T) { h := NewHub(5) var mu sync.Mutex registered := []string{} unregistered := []string{} h.SetPresenceHooks( func(userID, connID string) { mu.Lock() defer mu.Unlock() registered = append(registered, userID) }, func(userID, connID string) { mu.Lock() defer mu.Unlock() unregistered = append(unregistered, userID) }, ) conn := newTestConn(t) c, _ := h.Register("hook-user", conn) h.Unregister(c) if len(registered) != 1 || registered[0] != "hook-user" { t.Errorf("registered = %v, want [hook-user]", registered) } if len(unregistered) != 1 || unregistered[0] != "hook-user" { t.Errorf("unregistered = %v, want [hook-user]", unregistered) } } // TestForEachUser verifies iteration over online users (used by ISSUE-058 rebuild). func TestForEachUser(t *testing.T) { h := NewHub(5) conn1 := newTestConn(t) conn2 := newTestConn(t) c1, _ := h.Register("u-a", conn1) c2, _ := h.Register("u-b", conn2) defer h.Unregister(c1) defer h.Unregister(c2) seen := map[string]bool{} h.ForEachUser(func(userID string) { seen[userID] = true }) if len(seen) != 2 || !seen["u-a"] || !seen["u-b"] { t.Errorf("ForEachUser seen = %v, want {u-a, u-b}", seen) } } // TestConnectionSendOnClosed verifies Send returns false after Close. func TestConnectionSendOnClosed(t *testing.T) { h := NewHub(5) conn := newTestConn(t) c, _ := h.Register("u-close", conn) c.Close() if ok := c.Send([]byte("x")); ok { t.Error("Send after Close = true, want false") } } // TestConnectionSendChannelFull verifies Send returns false when channel is full. func TestConnectionSendChannelFull(t *testing.T) { h := NewHub(5) conn := newTestConn(t) c, _ := h.Register("u-full", conn) defer h.Unregister(c) // Fill the send channel (cap 64). for i := 0; i < sendBufferSize; i++ { if !c.Send([]byte("x")) { t.Fatalf("Send #%d returned false unexpectedly", i) } } // Next Send should drop. if ok := c.Send([]byte("overflow")); ok { t.Error("Send on full channel = true, want false") } } // TestNewHubDefaultMaxConns verifies NewHub applies a default when maxConns <= 0. func TestNewHubDefaultMaxConns(t *testing.T) { h := NewHub(0) if h.maxConnsPerUser != 5 { t.Errorf("NewHub(0) maxConnsPerUser = %d, want default 5", h.maxConnsPerUser) } h2 := NewHub(-1) if h2.maxConnsPerUser != 5 { t.Errorf("NewHub(-1) maxConnsPerUser = %d, want default 5", h2.maxConnsPerUser) } }