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") } // Verify the checker was invoked at least once if !called { t.Error("access checker was not called") } } // ---------- 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) client := &Client{WorkspaceID: "", Send: ch} h := NewHub(nil) h.clients = map[*Client]bool{client: 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) client := &Client{WorkspaceID: "ws-caller", Send: ch} 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{client: 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) { // Use a named variable for the client so the map key and Broadcast's // range both refer to the same *Client pointer. ch := make(chan []byte, 1) client := &Client{WorkspaceID: "", Send: ch} h := NewHub(nil) h.clients = map[*Client]bool{client: true} // Fill and close so any subsequent send (from Broadcast) hits // safeSend's default → returns false without blocking or panicking. ch <- []byte("fill") close(ch) // Broadcast must not panic — safeSend returns false for closed channel h.Broadcast(models.WSMessage{Type: "test"}) } 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 c1 := &Client{WorkspaceID: "ws-1", Send: ch1} c2 := &Client{WorkspaceID: "ws-2", Send: ch2} c3 := &Client{WorkspaceID: "ws-3", Send: ch3} 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{c1: true, c2: true, c3: 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) canvasClient := &Client{WorkspaceID: "", Send: ch} 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{ canvasClient: true, // canvas client &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") } }