pulse/internal/websocket/hub_tenant_test.go

package websocket

import (
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

// MockStateGetter implements StateGetter interface
type MockStateGetter struct {
	state interface{}
}

// MockTenantStateGetter implements TenantStateGetter interface
type MockTenantStateGetter struct {
	state map[string]interface{}
}

func (m *MockTenantStateGetter) GetState(orgID string) interface{} {
	return m.state[orgID]
}

// MockOrgAuthChecker implements OrgAuthChecker interface
type MockOrgAuthChecker struct {
	called bool
	allow  bool
}

func (m *MockOrgAuthChecker) CanAccessOrg(userID string, token interface{}, orgID string) bool {
	m.called = true
	return m.allow
}

func TestHub_Tenant_Broadcasting(t *testing.T) {
	hub := NewHub(nil)
	go hub.Run()
	defer func() {
		// Stop hub logic if exposed, or just let it die with test
	}()

	// Setup Tenant State Getter
	mockState := &MockTenantStateGetter{
		state: map[string]interface{}{
			"org1": map[string]string{"foo": "bar"},
			"org2": map[string]string{"baz": "qux"},
		},
	}
	hub.SetStateGetter(mockState.GetState)

	// Test GetTenantClientCount
	assert.Equal(t, 0, hub.GetTenantClientCount("org1"))

	// Test BroadcastStateToTenant (should not panic even with 0 clients)
	hub.BroadcastStateToTenant("org1", map[string]string{"status": "ok"})
	hub.BroadcastStateToTenant("org2", nil)
	hub.BroadcastStateToTenant("missing", nil)

	// Allow async broadcast to process
	time.Sleep(100 * time.Millisecond)
}

func TestHub_Setters_Coverage(t *testing.T) {
	hub := NewHub(nil)

	// Test SetOrgAuthChecker
	checker := &MockOrgAuthChecker{allow: true}
	hub.SetOrgAuthChecker(checker)

	// Verify it was set
	assert.NotNil(t, hub.orgAuthChecker)

	// Trigger the checker
	success := hub.orgAuthChecker.CanAccessOrg("user", "tok", "org")
	assert.True(t, success)
	assert.True(t, checker.called)
}

func TestHub_DispatchToTenantClients(t *testing.T) {
	// This tests the internal logic of iterating clients
	hub := NewHub(nil)
	go hub.Run()
	defer hub.Stop()

	// Create a mock client
	client := &Client{
		hub:   hub,
		send:  make(chan []byte, 256),
		orgID: "org1",
		// isActive removed
	}

	// Manually register (simulating register channel)
	hub.register <- client

	// Allow registration to process
	time.Sleep(50 * time.Millisecond)

	// Now broadcast to org1 (internal method)
	msg := []byte("test message")
	// dispatchToTenantClients(orgID string, data []byte, dropLog string)
	// But wait, dispatchToTenantClients is private (lowercase). Can we call it?
	// Tests are in package `websocket`, so yes we can access private methods of `hub.go`.
	hub.dispatchToTenantClients("org1", msg, "Dropping test message")

	// Check if client received it
	select {
	case received := <-client.send:
		assert.Equal(t, msg, received)
	case <-time.After(100 * time.Millisecond):
		t.Fatal("Client did not receive tenant broadcast")
	}

	// Broadcast to org2 (should not receive)
	hub.dispatchToTenantClients("org2", msg, "Dropping test message")
	select {
	case <-client.send:
		t.Fatal("Client received message for wrong tenant")
	case <-time.After(50 * time.Millisecond):
		// Success
	}
}

func TestHub_GlobalDropRemovesClientFromTenantIndex(t *testing.T) {
	hub := NewHub(nil)
	go hub.Run()
	t.Cleanup(hub.Stop)

	client := registerTenantTestClient(t, hub, "client-full", "org1")
	for i := 0; i < cap(client.send); i++ {
		client.send <- []byte("queued")
	}

	hub.dispatchToClients([]byte("drop"), "Dropping full client")

	hub.mu.RLock()
	_, inClients := hub.clients[client]
	tenantClients := hub.clientsByTenant["org1"]
	_, inTenant := tenantClients[client]
	hub.mu.RUnlock()

	if inClients {
		t.Fatal("full client remained in flat registry after global drop")
	}
	if inTenant || len(tenantClients) != 0 {
		t.Fatal("full client remained in tenant registry after global drop")
	}
	if !client.closed.Load() {
		t.Fatal("full client was not marked closed after global drop")
	}
}

func TestHub_TenantDropClearsStaleTenantOnlyClient(t *testing.T) {
	hub := NewHub(nil)
	client := &Client{
		hub:   hub,
		id:    "tenant-only-full",
		orgID: "org1",
		send:  make(chan []byte, 1),
	}
	client.send <- []byte("queued")

	hub.mu.Lock()
	hub.clientsByTenant["org1"] = map[*Client]bool{client: true}
	hub.mu.Unlock()

	hub.dispatchToTenantClients("org1", []byte("drop"), "Dropping stale tenant client")

	hub.mu.RLock()
	tenantClients := hub.clientsByTenant["org1"]
	_, inTenant := tenantClients[client]
	hub.mu.RUnlock()

	if inTenant || len(tenantClients) != 0 {
		t.Fatal("stale tenant-only client remained in tenant registry after tenant drop")
	}
	if !client.closed.Load() {
		t.Fatal("stale tenant-only client was not marked closed after tenant drop")
	}
}