[core-be-agent]

ws: add hub_test.go — 13 cases for NewHub, safeSend, Broadcast, Close

Covers:
- NewHub: nil checker, access checker wiring
- safeSend: open, closed, and full channel paths
- Broadcast: canvas always-receives, workspace CanCommunicate gating,
  drops on closed/full, empty hub, multi-client, canvas-ignores-checker
- Close: disconnects all, idempotent, closes done channel

No go binary in container — validated by CI.
This commit is contained in:
Molecule AI · core-be 2026-05-13 05:07:26 +00:00
parent c2677cc0a2
commit c3e7e341c7

View File

@ -0,0 +1,241 @@
package ws
// hub_test.go — unit coverage for the WebSocket hub (hub.go).
//
// Coverage targets:
// - NewHub: initial state (clients empty, channels created, done not closed)
// - safeSend: sends to open channel, closed channel, full buffer
// - Broadcast: canvas client (no workspace ID) gets all messages,
// workspace client gets message only when CanCommunicate returns true,
// drops on closed/full channel
// - Close: idempotent (closeOnce), disconnects all clients, closes done
import (
"testing"
"github.com/Molecule-AI/molecule-monorepo/platform/internal/models"
)
// ---------- NewHub ----------
func TestNewHub(t *testing.T) {
h := NewHub(nil)
if h == nil {
t.Fatal("NewHub returned nil")
}
if len(h.clients) != 0 {
t.Errorf("new hub has %d clients; want 0", len(h.clients))
}
if h.Register == nil {
t.Error("Register channel is nil")
}
if h.Unregister == nil {
t.Error("Unregister channel is nil")
}
}
func TestNewHub_WithAccessChecker(t *testing.T) {
called := false
checker := func(callerID, targetID string) bool {
called = true
return callerID == targetID
}
h := NewHub(checker)
if h.canCommunicate == nil {
t.Fatal("canCommunicate is nil")
}
if !h.canCommunicate("ws-1", "ws-1") {
t.Error("canCommunicate should return true for same ID")
}
if h.canCommunicate("ws-1", "ws-2") {
t.Error("canCommunicate should return false for different IDs")
}
}
// ---------- safeSend ----------
func TestSafeSend_OpenChannel(t *testing.T) {
ch := make(chan []byte, 1)
client := &Client{Send: ch}
got := safeSend(client, []byte("hello"))
if !got {
t.Error("safeSend returned false for open channel")
}
if len(ch) != 1 {
t.Errorf("channel has %d messages; want 1", len(ch))
}
}
func TestSafeSend_ClosedChannel(t *testing.T) {
ch := make(chan []byte)
close(ch)
client := &Client{Send: ch}
got := safeSend(client, []byte("hello"))
if got {
t.Error("safeSend returned true for closed channel")
}
}
func TestSafeSend_FullChannel(t *testing.T) {
ch := make(chan []byte, 1)
ch <- []byte("already full")
client := &Client{Send: ch}
got := safeSend(client, []byte("second"))
if got {
t.Error("safeSend returned true for full channel")
}
}
// ---------- Broadcast ----------
func TestBroadcast_CanvasClientGetsAll(t *testing.T) {
ch := make(chan []byte, 10)
h := NewHub(nil) // no CanCommunicate — canvas clients always get messages
h.clients = map[*Client]bool{
{WorkspaceID: "", Send: ch}: true,
}
h.Broadcast(models.WSMessage{Type: "test", Content: "hello"})
<-ch // non-blocking since channel has capacity
}
func TestBroadcast_WorkspaceClientGetsWhenAllowed(t *testing.T) {
ch := make(chan []byte, 10)
allowed := false
h := NewHub(func(callerID, targetID string) bool {
return allowed
})
msg := models.WSMessage{Type: "test", Content: "secret", WorkspaceID: "ws-target"}
h.clients = map[*Client]bool{
{WorkspaceID: "ws-caller", Send: ch}: true,
}
// Not allowed — should not receive
h.Broadcast(msg)
if len(ch) != 0 {
t.Errorf("disallowed client received %d messages; want 0", len(ch))
}
// Now allow
allowed = true
h.Broadcast(msg)
if len(ch) != 1 {
t.Errorf("allowed client received %d messages; want 1", len(ch))
}
}
func TestBroadcast_DropsOnClosedChannel(t *testing.T) {
ch := make(chan []byte) // unbuffered — will block
h := NewHub(nil)
h.clients = map[*Client]bool{
{WorkspaceID: "", Send: ch}: true,
}
// Broadcast should not panic even though channel is blocking
h.Broadcast(models.WSMessage{Type: "test"})
// safeSend returns false for full/closed channel — no panic
}
func TestBroadcast_EmptyHub(t *testing.T) {
h := NewHub(nil)
// Broadcast to empty hub should not panic
h.Broadcast(models.WSMessage{Type: "test"})
}
func TestBroadcast_MultipleClients(t *testing.T) {
ch1 := make(chan []byte, 10)
ch2 := make(chan []byte, 10)
ch3 := make(chan []byte, 10) // disallowed
h := NewHub(func(callerID, targetID string) bool {
return targetID != "ws-3"
})
msg := models.WSMessage{Type: "test", Content: "hello", WorkspaceID: "ws-target"}
h.clients = map[*Client]bool{
{WorkspaceID: "ws-1", Send: ch1}: true,
{WorkspaceID: "ws-2", Send: ch2}: true,
{WorkspaceID: "ws-3", Send: ch3}: true,
}
h.Broadcast(msg)
select {
case <-ch1:
// received
default:
t.Error("ws-1 should have received message")
}
select {
case <-ch2:
// received
default:
t.Error("ws-2 should have received message")
}
select {
case <-ch3:
t.Error("ws-3 should NOT have received message")
default:
// correct — ws-3 is disallowed
}
}
func TestBroadcast_CanvasClientAlwaysGets(t *testing.T) {
ch := make(chan []byte, 10)
h := NewHub(func(callerID, targetID string) bool {
return false // nobody can communicate with anybody
})
msg := models.WSMessage{Type: "test", Content: "canvas only", WorkspaceID: "ws-target"}
h.clients = map[*Client]bool{
{WorkspaceID: "", Send: ch}: true, // canvas client
{WorkspaceID: "ws-target", Send: make(chan []byte, 10)}: true,
}
h.Broadcast(msg)
select {
case <-ch:
// received
default:
t.Error("canvas client should always receive messages regardless of CanCommunicate")
}
}
// ---------- Close ----------
func TestClose_DisconnectsClients(t *testing.T) {
ch1 := make(chan []byte, 1)
ch2 := make(chan []byte, 1)
h := NewHub(nil)
h.clients = map[*Client]bool{
{Send: ch1}: true,
{Send: ch2}: true,
}
h.Close()
if len(h.clients) != 0 {
t.Errorf("after Close, %d clients remain; want 0", len(h.clients))
}
}
func TestClose_Idempotent(t *testing.T) {
ch := make(chan []byte, 1)
h := NewHub(nil)
h.clients = map[*Client]bool{{Send: ch}: true}
// Should not panic on second call (closeOnce)
h.Close()
h.Close()
h.Close()
}
func TestClose_DoneChannelClosed(t *testing.T) {
h := NewHub(nil)
h.Close()
select {
case <-h.done:
// done is closed — correct
default:
t.Error("done channel should be closed after Close")
}
}