molecule-core/workspace-server/internal/ws/hub_test.go

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