516 lines
13 KiB
Go
516 lines
13 KiB
Go
|
package gateway
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"fmt"
|
||
|
|
"log/slog"
|
||
|
|
"sync"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"code.cubecraftcreations.com/CubeCraft-Creations/Control-Center/go-backend/internal/handler"
|
||
|
|
"code.cubecraftcreations.com/CubeCraft-Creations/Control-Center/go-backend/internal/models"
|
||
|
|
)
|
||
|
|
|
||
|
|
// ── Mock AgentRepo ────────────────────────────────────────────────────────
|
||
|
|
|
||
|
|
type mockAgentRepo struct {
|
||
|
|
mu sync.Mutex
|
||
|
|
agents map[string]models.AgentCardData
|
||
|
|
updateCalls []updateCall
|
||
|
|
}
|
||
|
|
|
||
|
|
type updateCall struct {
|
||
|
|
id string
|
||
|
|
req models.UpdateAgentRequest
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *mockAgentRepo) Get(_ context.Context, id string) (models.AgentCardData, error) {
|
||
|
|
m.mu.Lock()
|
||
|
|
defer m.mu.Unlock()
|
||
|
|
a, ok := m.agents[id]
|
||
|
|
if !ok {
|
||
|
|
return models.AgentCardData{}, errNotFound
|
||
|
|
}
|
||
|
|
return a, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *mockAgentRepo) Update(_ context.Context, id string, req models.UpdateAgentRequest) (models.AgentCardData, error) {
|
||
|
|
m.mu.Lock()
|
||
|
|
defer m.mu.Unlock()
|
||
|
|
|
||
|
|
a, ok := m.agents[id]
|
||
|
|
if !ok {
|
||
|
|
return models.AgentCardData{}, errNotFound
|
||
|
|
}
|
||
|
|
|
||
|
|
if req.Status != nil {
|
||
|
|
a.Status = *req.Status
|
||
|
|
}
|
||
|
|
if req.DisplayName != nil {
|
||
|
|
a.DisplayName = *req.DisplayName
|
||
|
|
}
|
||
|
|
if req.Role != nil {
|
||
|
|
a.Role = *req.Role
|
||
|
|
}
|
||
|
|
if req.Channel != nil {
|
||
|
|
a.Channel = *req.Channel
|
||
|
|
}
|
||
|
|
if req.CurrentTask != nil {
|
||
|
|
a.CurrentTask = req.CurrentTask
|
||
|
|
}
|
||
|
|
if req.TaskProgress != nil {
|
||
|
|
a.TaskProgress = req.TaskProgress
|
||
|
|
}
|
||
|
|
if req.TaskElapsed != nil {
|
||
|
|
a.TaskElapsed = req.TaskElapsed
|
||
|
|
}
|
||
|
|
if req.ErrorMessage != nil {
|
||
|
|
a.ErrorMessage = req.ErrorMessage
|
||
|
|
}
|
||
|
|
if req.LastActivityAt != nil {
|
||
|
|
a.LastActivity = *req.LastActivityAt
|
||
|
|
}
|
||
|
|
|
||
|
|
m.agents[id] = a
|
||
|
|
m.updateCalls = append(m.updateCalls, updateCall{id, req})
|
||
|
|
return a, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *mockAgentRepo) Create(_ context.Context, a models.AgentCardData) error {
|
||
|
|
m.mu.Lock()
|
||
|
|
defer m.mu.Unlock()
|
||
|
|
m.agents[a.ID] = a
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *mockAgentRepo) List(_ context.Context, statusFilter models.AgentStatus) ([]models.AgentCardData, error) {
|
||
|
|
m.mu.Lock()
|
||
|
|
defer m.mu.Unlock()
|
||
|
|
|
||
|
|
var result []models.AgentCardData
|
||
|
|
for _, a := range m.agents {
|
||
|
|
if statusFilter == "" || a.Status == statusFilter {
|
||
|
|
result = append(result, a)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return result, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *mockAgentRepo) Delete(_ context.Context, id string) error {
|
||
|
|
m.mu.Lock()
|
||
|
|
defer m.mu.Unlock()
|
||
|
|
delete(m.agents, id)
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *mockAgentRepo) Count(_ context.Context) (int, error) {
|
||
|
|
m.mu.Lock()
|
||
|
|
defer m.mu.Unlock()
|
||
|
|
return len(m.agents), nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// errNotFound is returned by the mock repo when an agent is not found.
|
||
|
|
var errNotFound = fmt.Errorf("not found")
|
||
|
|
|
||
|
|
// ── Broadcast capture helper ───────────────────────────────────────────────
|
||
|
|
|
||
|
|
// broadcastCapture wraps a real Broker and captures all broadcasts
|
||
|
|
// via a subscribed channel. Use captured() to retrieve events that have
|
||
|
|
// been received so far. Call close() to unsubscribe when done.
|
||
|
|
type broadcastCapture struct {
|
||
|
|
broker *handler.Broker
|
||
|
|
ch chan handler.SSEEvent
|
||
|
|
}
|
||
|
|
|
||
|
|
func newBroadcastCapture(broker *handler.Broker) *broadcastCapture {
|
||
|
|
return &broadcastCapture{
|
||
|
|
broker: broker,
|
||
|
|
ch: broker.Subscribe(),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// captured drains all pending events from the subscription channel
|
||
|
|
// and returns them. This is synchronous — it only returns events that
|
||
|
|
// have already been sent to the channel.
|
||
|
|
func (bc *broadcastCapture) captured() []handler.SSEEvent {
|
||
|
|
var events []handler.SSEEvent
|
||
|
|
for {
|
||
|
|
select {
|
||
|
|
case evt := <-bc.ch:
|
||
|
|
events = append(events, evt)
|
||
|
|
default:
|
||
|
|
return events
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (bc *broadcastCapture) close() {
|
||
|
|
bc.broker.Unsubscribe(bc.ch)
|
||
|
|
}
|
||
|
|
|
||
|
|
// ── Test helpers ──────────────────────────────────────────────────────────
|
||
|
|
|
||
|
|
// newTestWSClient creates a WSClient wired to a mock repo and a real broker.
|
||
|
|
// Returns the client, the mock repo, and a broadcast capture.
|
||
|
|
func newTestWSClient() (*WSClient, *mockAgentRepo, *handler.Broker, *broadcastCapture) {
|
||
|
|
repo := &mockAgentRepo{agents: make(map[string]models.AgentCardData)}
|
||
|
|
broker := handler.NewBroker()
|
||
|
|
capture := newBroadcastCapture(broker)
|
||
|
|
client := NewWSClient(WSConfig{}, repo, broker, slog.Default())
|
||
|
|
return client, repo, broker, capture
|
||
|
|
}
|
||
|
|
|
||
|
|
// ── Tests ─────────────────────────────────────────────────────────────────
|
||
|
|
|
||
|
|
func TestHandleSessionsChanged_Active(t *testing.T) {
|
||
|
|
client, repo, broker, capture := newTestWSClient()
|
||
|
|
defer capture.close()
|
||
|
|
|
||
|
|
repo.agents["otto"] = models.AgentCardData{
|
||
|
|
ID: "otto",
|
||
|
|
DisplayName: "Otto",
|
||
|
|
Status: models.AgentStatusIdle,
|
||
|
|
}
|
||
|
|
|
||
|
|
payload := json.RawMessage(`{
|
||
|
|
"sessionKey": "s1",
|
||
|
|
"agentId": "otto",
|
||
|
|
"status": "running",
|
||
|
|
"totalTokens": 500,
|
||
|
|
"currentTask": "Orchestrating tasks"
|
||
|
|
}`)
|
||
|
|
|
||
|
|
client.handleSessionsChanged(payload)
|
||
|
|
|
||
|
|
// Verify: agent status updated to active
|
||
|
|
repo.mu.Lock()
|
||
|
|
agent := repo.agents["otto"]
|
||
|
|
calls := make([]updateCall, len(repo.updateCalls))
|
||
|
|
copy(calls, repo.updateCalls)
|
||
|
|
repo.mu.Unlock()
|
||
|
|
|
||
|
|
if agent.Status != models.AgentStatusActive {
|
||
|
|
t.Errorf("agent status = %q, want %q", agent.Status, models.AgentStatusActive)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Verify: update was called
|
||
|
|
if len(calls) == 0 {
|
||
|
|
t.Fatal("expected at least one update call")
|
||
|
|
}
|
||
|
|
if calls[0].id != "otto" {
|
||
|
|
t.Errorf("update call agentId = %q, want %q", calls[0].id, "otto")
|
||
|
|
}
|
||
|
|
|
||
|
|
// Verify: broker broadcast "agent.status"
|
||
|
|
events := capture.captured()
|
||
|
|
found := false
|
||
|
|
for _, evt := range events {
|
||
|
|
if evt.EventType == "agent.status" {
|
||
|
|
found = true
|
||
|
|
break
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if !found {
|
||
|
|
t.Error("expected broker broadcast with event type 'agent.status'")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestHandleSessionsChanged_Idle(t *testing.T) {
|
||
|
|
client, repo, _, capture := newTestWSClient()
|
||
|
|
defer capture.close()
|
||
|
|
|
||
|
|
repo.agents["dex"] = models.AgentCardData{
|
||
|
|
ID: "dex",
|
||
|
|
DisplayName: "Dex",
|
||
|
|
Status: models.AgentStatusActive,
|
||
|
|
CurrentTask: strPtr("Writing API"),
|
||
|
|
}
|
||
|
|
|
||
|
|
payload := json.RawMessage(`{
|
||
|
|
"sessionKey": "s2",
|
||
|
|
"agentId": "dex",
|
||
|
|
"status": "done",
|
||
|
|
"totalTokens": 1000
|
||
|
|
}`)
|
||
|
|
|
||
|
|
client.handleSessionsChanged(payload)
|
||
|
|
|
||
|
|
repo.mu.Lock()
|
||
|
|
agent := repo.agents["dex"]
|
||
|
|
repo.mu.Unlock()
|
||
|
|
|
||
|
|
// Verify: agent goes idle
|
||
|
|
if agent.Status != models.AgentStatusIdle {
|
||
|
|
t.Errorf("agent status = %q, want %q", agent.Status, models.AgentStatusIdle)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Verify: current task cleared (set to empty string)
|
||
|
|
if agent.CurrentTask != nil && *agent.CurrentTask != "" {
|
||
|
|
t.Errorf("current task = %q, want empty (cleared on idle)", *agent.CurrentTask)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Verify: broker fires "agent.status"
|
||
|
|
events := capture.captured()
|
||
|
|
found := false
|
||
|
|
for _, evt := range events {
|
||
|
|
if evt.EventType == "agent.status" {
|
||
|
|
found = true
|
||
|
|
break
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if !found {
|
||
|
|
t.Error("expected broker broadcast with event type 'agent.status'")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestHandleSessionsChanged_ArrayPayload(t *testing.T) {
|
||
|
|
client, repo, _, capture := newTestWSClient()
|
||
|
|
defer capture.close()
|
||
|
|
|
||
|
|
repo.agents["otto"] = models.AgentCardData{ID: "otto", DisplayName: "Otto", Status: models.AgentStatusIdle}
|
||
|
|
repo.agents["dex"] = models.AgentCardData{ID: "dex", DisplayName: "Dex", Status: models.AgentStatusIdle}
|
||
|
|
|
||
|
|
payload := json.RawMessage(`[
|
||
|
|
{"sessionKey":"s1","agentId":"otto","status":"running","totalTokens":100},
|
||
|
|
{"sessionKey":"s2","agentId":"dex","status":"streaming","totalTokens":200}
|
||
|
|
]`)
|
||
|
|
|
||
|
|
client.handleSessionsChanged(payload)
|
||
|
|
|
||
|
|
repo.mu.Lock()
|
||
|
|
otto := repo.agents["otto"]
|
||
|
|
dex := repo.agents["dex"]
|
||
|
|
repo.mu.Unlock()
|
||
|
|
|
||
|
|
if otto.Status != models.AgentStatusActive {
|
||
|
|
t.Errorf("otto status = %q, want active", otto.Status)
|
||
|
|
}
|
||
|
|
if dex.Status != models.AgentStatusActive {
|
||
|
|
t.Errorf("dex status = %q, want active", dex.Status)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Both should produce broadcasts
|
||
|
|
events := capture.captured()
|
||
|
|
statusCount := 0
|
||
|
|
for _, evt := range events {
|
||
|
|
if evt.EventType == "agent.status" {
|
||
|
|
statusCount++
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if statusCount < 2 {
|
||
|
|
t.Errorf("expected at least 2 agent.status broadcasts, got %d", statusCount)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestHandleSessionsChanged_SkipsEmptyAgentID(t *testing.T) {
|
||
|
|
client, _, _, capture := newTestWSClient()
|
||
|
|
defer capture.close()
|
||
|
|
|
||
|
|
payload := json.RawMessage(`{"sessionKey":"s1","agentId":"","status":"running"}`)
|
||
|
|
client.handleSessionsChanged(payload)
|
||
|
|
|
||
|
|
events := capture.captured()
|
||
|
|
if len(events) > 0 {
|
||
|
|
t.Errorf("expected no broadcasts for empty agentId, got %d", len(events))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestHandleSessionsChanged_UnparseablePayload(t *testing.T) {
|
||
|
|
client, _, _, capture := newTestWSClient()
|
||
|
|
defer capture.close()
|
||
|
|
|
||
|
|
payload := json.RawMessage(`not json at all`)
|
||
|
|
client.handleSessionsChanged(payload)
|
||
|
|
|
||
|
|
events := capture.captured()
|
||
|
|
if len(events) > 0 {
|
||
|
|
t.Errorf("expected no broadcasts for unparseable payload, got %d", len(events))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestHandlePresence(t *testing.T) {
|
||
|
|
client, repo, _, capture := newTestWSClient()
|
||
|
|
defer capture.close()
|
||
|
|
|
||
|
|
repo.agents["pip"] = models.AgentCardData{
|
||
|
|
ID: "pip",
|
||
|
|
DisplayName: "Pip",
|
||
|
|
Status: models.AgentStatusActive,
|
||
|
|
}
|
||
|
|
|
||
|
|
payload := json.RawMessage(`{
|
||
|
|
"agentId": "pip",
|
||
|
|
"connected": true,
|
||
|
|
"lastActivityAt": "2025-01-01T00:00:00Z"
|
||
|
|
}`)
|
||
|
|
|
||
|
|
client.handlePresence(payload)
|
||
|
|
|
||
|
|
repo.mu.Lock()
|
||
|
|
agent := repo.agents["pip"]
|
||
|
|
calls := make([]updateCall, len(repo.updateCalls))
|
||
|
|
copy(calls, repo.updateCalls)
|
||
|
|
repo.mu.Unlock()
|
||
|
|
|
||
|
|
// Agent should still be active (connected=true doesn't change status)
|
||
|
|
if agent.Status != models.AgentStatusActive {
|
||
|
|
t.Errorf("agent status = %q, want active", agent.Status)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Update should have been called (for lastActivityAt)
|
||
|
|
if len(calls) == 0 {
|
||
|
|
t.Fatal("expected at least one update call")
|
||
|
|
}
|
||
|
|
|
||
|
|
// Verify broadcast
|
||
|
|
events := capture.captured()
|
||
|
|
found := false
|
||
|
|
for _, evt := range events {
|
||
|
|
if evt.EventType == "agent.status" {
|
||
|
|
found = true
|
||
|
|
break
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if !found {
|
||
|
|
t.Error("expected broker broadcast with event type 'agent.status'")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestHandlePresence_Disconnect(t *testing.T) {
|
||
|
|
client, repo, _, capture := newTestWSClient()
|
||
|
|
defer capture.close()
|
||
|
|
|
||
|
|
repo.agents["pip"] = models.AgentCardData{
|
||
|
|
ID: "pip",
|
||
|
|
DisplayName: "Pip",
|
||
|
|
Status: models.AgentStatusActive,
|
||
|
|
}
|
||
|
|
|
||
|
|
payload := json.RawMessage(`{
|
||
|
|
"agentId": "pip",
|
||
|
|
"connected": false
|
||
|
|
}`)
|
||
|
|
|
||
|
|
client.handlePresence(payload)
|
||
|
|
|
||
|
|
repo.mu.Lock()
|
||
|
|
agent := repo.agents["pip"]
|
||
|
|
repo.mu.Unlock()
|
||
|
|
|
||
|
|
// Agent should go idle on disconnect
|
||
|
|
if agent.Status != models.AgentStatusIdle {
|
||
|
|
t.Errorf("agent status = %q, want idle after disconnect", agent.Status)
|
||
|
|
}
|
||
|
|
|
||
|
|
events := capture.captured()
|
||
|
|
found := false
|
||
|
|
for _, evt := range events {
|
||
|
|
if evt.EventType == "agent.status" {
|
||
|
|
found = true
|
||
|
|
break
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if !found {
|
||
|
|
t.Error("expected broker broadcast with event type 'agent.status' on disconnect")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestHandlePresence_EmptyAgentID(t *testing.T) {
|
||
|
|
client, _, _, capture := newTestWSClient()
|
||
|
|
defer capture.close()
|
||
|
|
|
||
|
|
payload := json.RawMessage(`{"agentId":"","connected":true}`)
|
||
|
|
client.handlePresence(payload)
|
||
|
|
|
||
|
|
events := capture.captured()
|
||
|
|
if len(events) > 0 {
|
||
|
|
t.Errorf("expected no broadcasts for empty agentId, got %d", len(events))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestHandleAgentConfig(t *testing.T) {
|
||
|
|
client, repo, _, capture := newTestWSClient()
|
||
|
|
defer capture.close()
|
||
|
|
|
||
|
|
repo.agents["rex"] = models.AgentCardData{
|
||
|
|
ID: "rex",
|
||
|
|
DisplayName: "Rex",
|
||
|
|
Role: "Frontend Dev",
|
||
|
|
Status: models.AgentStatusIdle,
|
||
|
|
Channel: "discord",
|
||
|
|
}
|
||
|
|
|
||
|
|
payload := json.RawMessage(`{
|
||
|
|
"id": "rex",
|
||
|
|
"name": "Rex the Dev",
|
||
|
|
"role": "Senior Frontend",
|
||
|
|
"channel": "telegram"
|
||
|
|
}`)
|
||
|
|
|
||
|
|
client.handleAgentConfig(payload)
|
||
|
|
|
||
|
|
repo.mu.Lock()
|
||
|
|
agent := repo.agents["rex"]
|
||
|
|
calls := make([]updateCall, len(repo.updateCalls))
|
||
|
|
copy(calls, repo.updateCalls)
|
||
|
|
repo.mu.Unlock()
|
||
|
|
|
||
|
|
// Verify DisplayName and Role updated
|
||
|
|
if agent.DisplayName != "Rex the Dev" {
|
||
|
|
t.Errorf("displayName = %q, want %q", agent.DisplayName, "Rex the Dev")
|
||
|
|
}
|
||
|
|
if agent.Role != "Senior Frontend" {
|
||
|
|
t.Errorf("role = %q, want %q", agent.Role, "Senior Frontend")
|
||
|
|
}
|
||
|
|
if agent.Channel != "telegram" {
|
||
|
|
t.Errorf("channel = %q, want %q", agent.Channel, "telegram")
|
||
|
|
}
|
||
|
|
|
||
|
|
// Verify update was called
|
||
|
|
if len(calls) == 0 {
|
||
|
|
t.Fatal("expected at least one update call")
|
||
|
|
}
|
||
|
|
|
||
|
|
// Verify broker fires "fleet.update"
|
||
|
|
events := capture.captured()
|
||
|
|
found := false
|
||
|
|
for _, evt := range events {
|
||
|
|
if evt.EventType == "fleet.update" {
|
||
|
|
found = true
|
||
|
|
break
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if !found {
|
||
|
|
t.Error("expected broker broadcast with event type 'fleet.update'")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestHandleAgentConfig_EmptyID(t *testing.T) {
|
||
|
|
client, _, _, capture := newTestWSClient()
|
||
|
|
defer capture.close()
|
||
|
|
|
||
|
|
payload := json.RawMessage(`{"id":"","name":"Ghost"}`)
|
||
|
|
client.handleAgentConfig(payload)
|
||
|
|
|
||
|
|
events := capture.captured()
|
||
|
|
if len(events) > 0 {
|
||
|
|
t.Errorf("expected no broadcasts for empty id, got %d", len(events))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestHandleAgentConfig_NotFound(t *testing.T) {
|
||
|
|
client, _, _, capture := newTestWSClient()
|
||
|
|
defer capture.close()
|
||
|
|
|
||
|
|
payload := json.RawMessage(`{"id":"unknown","name":"Ghost","role":"Phantom"}`)
|
||
|
|
client.handleAgentConfig(payload)
|
||
|
|
|
||
|
|
// Agent doesn't exist in repo, so Update will fail → handler logs warning, returns early
|
||
|
|
events := capture.captured()
|
||
|
|
for _, evt := range events {
|
||
|
|
if evt.EventType == "fleet.update" {
|
||
|
|
t.Error("fleet.update should not be broadcast when agent update fails")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|