Edu/services/push-gateway/internal/hub/hub_test.go

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)
	}
}