arp/arp_agent/agent_test.go

package main

import (
	"context"
	"encoding/json"
	"strings"
	"testing"
	"time"

	"github.com/sashabaranov/go-openai"
)

// TestAgent_Initialize tests agent initialization
func TestAgent_Initialize(t *testing.T) {
	mockMCP := NewMockMCPClient([]Tool{
		{
			Name:        "introspect",
			Description: "Discover the GraphQL schema",
			InputSchema: InputSchema{
				Type:                 "object",
				Properties:           map[string]Property{},
				AdditionalProperties: false,
			},
		},
		{
			Name:        "query",
			Description: "Execute a GraphQL query",
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"query": {Type: "string", Description: "The GraphQL query"},
				},
				Required:             []string{"query"},
				AdditionalProperties: false,
			},
		},
	})

	mockLLM := NewMockLLM(nil)
	agent := NewTestAgent(mockLLM, mockMCP)

	err := agent.Initialize()
	if err != nil {
		t.Fatalf("Initialize failed: %v", err)
	}

	if len(agent.tools) != 2 {
		t.Errorf("Expected 2 tools, got %d", len(agent.tools))
	}

	// Verify tools were converted correctly
	toolNames := make([]string, len(agent.tools))
	for i, tool := range agent.tools {
		toolNames[i] = tool.Function.Name
	}

	expectedNames := []string{"introspect", "query"}
	for i, expected := range expectedNames {
		if toolNames[i] != expected {
			t.Errorf("Tool %d: expected name %s, got %s", i, expected, toolNames[i])
		}
	}
}

// TestAgent_ProcessEvent tests event processing
func TestAgent_ProcessEvent(t *testing.T) {
	ctx := context.Background()

	t.Run("TaskCreatedEvent", func(t *testing.T) {
		mockMCP := NewMockMCPClient([]Tool{
			{Name: "query", Description: "Execute a GraphQL query", InputSchema: InputSchema{Type: "object"}},
		})
		mockMCP.SetToolResult("query", &CallToolResult{
			Content: []ContentBlock{{Type: "text", Text: `{"data": {"tasks": []}}`}},
		})

		// Mock LLM that makes a tool call then responds
		mockLLM := NewMockLLM([]*openai.ChatCompletionMessage{
			{
				Role: openai.ChatMessageRoleAssistant,
				ToolCalls: []openai.ToolCall{
					{
						ID: "call-1",
						Function: openai.FunctionCall{
							Name:      "query",
							Arguments: `{"query": "{ tasks { id title } }"}`,
						},
					},
				},
			},
			{
				Role:    openai.ChatMessageRoleAssistant,
				Content: "I've processed the task created event.",
			},
		})

		agent := NewTestAgent(mockLLM, mockMCP)
		agent.Initialize()

		eventData := json.RawMessage(`{"taskId": "task-123", "title": "New Task"}`)
		err := agent.ProcessEvent(ctx, "graphql://subscription/taskCreated", eventData)
		if err != nil {
			t.Errorf("ProcessEvent failed: %v", err)
		}
	})

	t.Run("MessageAddedEvent", func(t *testing.T) {
		mockMCP := NewMockMCPClient([]Tool{
			{Name: "query", Description: "Execute a GraphQL query", InputSchema: InputSchema{Type: "object"}},
		})

		// Mock LLM that responds directly without tool calls
		mockLLM := NewMockLLM([]*openai.ChatCompletionMessage{
			{
				Role:    openai.ChatMessageRoleAssistant,
				Content: "I received the message added event.",
			},
		})

		agent := NewTestAgent(mockLLM, mockMCP)
		agent.Initialize()

		eventData := json.RawMessage(`{"messageId": "msg-456", "content": "Hello!"}`)
		err := agent.ProcessEvent(ctx, "graphql://subscription/messageAdded", eventData)
		if err != nil {
			t.Errorf("ProcessEvent failed: %v", err)
		}
	})
}

// TestAgent_Run tests the interactive Run method
func TestAgent_Run(t *testing.T) {
	ctx := context.Background()

	t.Run("SimpleResponse", func(t *testing.T) {
		mockMCP := NewMockMCPClient([]Tool{})
		mockLLM := NewMockLLM([]*openai.ChatCompletionMessage{
			{
				Role:    openai.ChatMessageRoleAssistant,
				Content: "Hello! How can I help you?",
			},
		})

		agent := NewTestAgent(mockLLM, mockMCP)
		agent.Initialize()

		response, err := agent.Run(ctx, "Hello")
		if err != nil {
			t.Fatalf("Run failed: %v", err)
		}

		if response != "Hello! How can I help you?" {
			t.Errorf("Expected 'Hello! How can I help you?', got '%s'", response)
		}
	})

	t.Run("WithToolCall", func(t *testing.T) {
		mockMCP := NewMockMCPClient([]Tool{
			{Name: "introspect", Description: "Introspect schema", InputSchema: InputSchema{Type: "object"}},
		})
		mockMCP.SetToolResult("introspect", &CallToolResult{
			Content: []ContentBlock{{Type: "text", Text: "Schema: Query, Mutation, Subscription"}},
		})

		mockLLM := NewMockLLM([]*openai.ChatCompletionMessage{
			{
				Role: openai.ChatMessageRoleAssistant,
				ToolCalls: []openai.ToolCall{
					{
						ID: "call-1",
						Function: openai.FunctionCall{
							Name:      "introspect",
							Arguments: `{}`,
						},
					},
				},
			},
			{
				Role:    openai.ChatMessageRoleAssistant,
				Content: "The schema has Query, Mutation, and Subscription types.",
			},
		})

		agent := NewTestAgent(mockLLM, mockMCP)
		agent.Initialize()

		response, err := agent.Run(ctx, "What types are in the schema?")
		if err != nil {
			t.Fatalf("Run failed: %v", err)
		}

		if response != "The schema has Query, Mutation, and Subscription types." {
			t.Errorf("Unexpected response: %s", response)
		}
	})
}

// TestAgent_BuildEventPrompt tests event prompt building
func TestAgent_BuildEventPrompt(t *testing.T) {
	tests := []struct {
		name      string
		uri       string
		eventData json.RawMessage
		wantType  string
	}{
		{
			name:      "TaskCreated",
			uri:       "graphql://subscription/taskCreated",
			eventData: json.RawMessage(`{"id": "1"}`),
			wantType:  "task created",
		},
		{
			name:      "TaskUpdated",
			uri:       "graphql://subscription/taskUpdated",
			eventData: json.RawMessage(`{"id": "2"}`),
			wantType:  "task updated",
		},
		{
			name:      "TaskDeleted",
			uri:       "graphql://subscription/taskDeleted",
			eventData: json.RawMessage(`{"id": "3"}`),
			wantType:  "task deleted",
		},
		{
			name:      "MessageAdded",
			uri:       "graphql://subscription/messageAdded",
			eventData: json.RawMessage(`{"id": "4"}`),
			wantType:  "message added",
		},
		{
			name:      "UnknownEvent",
			uri:       "graphql://subscription/unknown",
			eventData: json.RawMessage(`{}`),
			wantType:  "unknown",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			prompt := buildTestEventPrompt(tt.uri, tt.eventData)
			if !strings.Contains(prompt, tt.wantType) {
				t.Errorf("Expected prompt to contain '%s', got: %s", tt.wantType, prompt)
			}
			if !strings.Contains(prompt, tt.uri) {
				t.Errorf("Expected prompt to contain URI '%s', got: %s", tt.uri, prompt)
			}
		})
	}
}

// TestAgent_ToolNames tests the toolNames helper function
func TestAgent_ToolNames(t *testing.T) {
	tools := []Tool{
		{Name: "introspect"},
		{Name: "query"},
		{Name: "mutate"},
	}

	names := toolNames(tools)

	expected := []string{"introspect", "query", "mutate"}
	if len(names) != len(expected) {
		t.Errorf("Expected %d names, got %d", len(expected), len(names))
	}

	for i, name := range names {
		if name != expected[i] {
			t.Errorf("Name %d: expected %s, got %s", i, expected[i], name)
		}
	}
}

// TestEventQueue tests the EventQueue operations
func TestEventQueue(t *testing.T) {
	t.Run("TryEnqueueSuccess", func(t *testing.T) {
		queue := NewEventQueue("test", 10)
		event := &QueuedEvent{
			URI:       "test://uri",
			Data:      json.RawMessage(`{"test": "data"}`),
			Timestamp: time.Now(),
		}

		success := queue.TryEnqueue(event)
		if !success {
			t.Error("Expected TryEnqueue to succeed")
		}
		if queue.Len() != 1 {
			t.Errorf("Expected queue length 1, got %d", queue.Len())
		}
	})

	t.Run("TryEnqueueFullQueue", func(t *testing.T) {
		queue := NewEventQueue("test", 2)

		// Fill the queue
		for i := 0; i < 2; i++ {
			success := queue.TryEnqueue(&QueuedEvent{URI: "test://uri"})
			if !success {
				t.Errorf("Expected TryEnqueue %d to succeed", i)
			}
		}

		// Try to add one more - should fail
		success := queue.TryEnqueue(&QueuedEvent{URI: "test://overflow"})
		if success {
			t.Error("Expected TryEnqueue to fail on full queue")
		}
		if queue.Len() != 2 {
			t.Errorf("Expected queue length 2, got %d", queue.Len())
		}
	})

	t.Run("Dequeue", func(t *testing.T) {
		queue := NewEventQueue("test", 10)
		event1 := &QueuedEvent{URI: "test://uri1"}
		event2 := &QueuedEvent{URI: "test://uri2"}

		queue.TryEnqueue(event1)
		queue.TryEnqueue(event2)

		// Dequeue should return events in FIFO order
		dequeued1 := queue.Dequeue()
		if dequeued1.URI != "test://uri1" {
			t.Errorf("Expected URI 'test://uri1', got '%s'", dequeued1.URI)
		}

		dequeued2 := queue.Dequeue()
		if dequeued2.URI != "test://uri2" {
			t.Errorf("Expected URI 'test://uri2', got '%s'", dequeued2.URI)
		}
	})

	t.Run("Len", func(t *testing.T) {
		queue := NewEventQueue("test", 10)

		if queue.Len() != 0 {
			t.Errorf("Expected empty queue to have length 0, got %d", queue.Len())
		}

		queue.TryEnqueue(&QueuedEvent{URI: "test://uri1"})
		if queue.Len() != 1 {
			t.Errorf("Expected queue length 1, got %d", queue.Len())
		}

		queue.TryEnqueue(&QueuedEvent{URI: "test://uri2"})
		if queue.Len() != 2 {
			t.Errorf("Expected queue length 2, got %d", queue.Len())
		}
	})
}

// TestAgent_QueueEvent tests event routing to queues
func TestAgent_QueueEvent(t *testing.T) {
	mockMCP := NewMockMCPClient([]Tool{})
	mockLLM := NewMockLLM(nil)
	agent := NewTestAgent(mockLLM, mockMCP)
	agent.SetupQueues(10)

	tests := []struct {
		name         string
		uri          string
		expectedTask int
		expectedMsg  int
	}{
		{
			name:         "TaskCreated",
			uri:          "graphql://subscription/taskCreated",
			expectedTask: 1,
			expectedMsg:  0,
		},
		{
			name:         "TaskUpdated",
			uri:          "graphql://subscription/taskUpdated",
			expectedTask: 1,
			expectedMsg:  0,
		},
		{
			name:         "TaskDeleted",
			uri:          "graphql://subscription/taskDeleted",
			expectedTask: 1,
			expectedMsg:  0,
		},
		{
			name:         "MessageAdded",
			uri:          "graphql://subscription/messageAdded",
			expectedTask: 0,
			expectedMsg:  1,
		},
		{
			name:         "UnknownEvent",
			uri:          "graphql://subscription/unknown",
			expectedTask: 1, // Unknown events go to task queue
			expectedMsg:  0,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Reset queues
			agent.SetupQueues(10)

			agent.QueueEvent(tt.uri, json.RawMessage(`{}`))

			stats := agent.GetQueueStats()
			if stats.TaskQueueSize != tt.expectedTask {
				t.Errorf("Expected task queue size %d, got %d", tt.expectedTask, stats.TaskQueueSize)
			}
			if stats.MessageQueueSize != tt.expectedMsg {
				t.Errorf("Expected message queue size %d, got %d", tt.expectedMsg, stats.MessageQueueSize)
			}
		})
	}
}

// TestAgent_QueueEventFullQueue tests that events are dropped when queue is full
func TestAgent_QueueEventFullQueue(t *testing.T) {
	mockMCP := NewMockMCPClient([]Tool{})
	mockLLM := NewMockLLM(nil)
	agent := NewTestAgent(mockLLM, mockMCP)
	agent.SetupQueues(2) // Small queue for testing

	// Fill the task queue
	agent.QueueEvent("graphql://subscription/taskCreated", json.RawMessage(`{"id": "1"}`))
	agent.QueueEvent("graphql://subscription/taskCreated", json.RawMessage(`{"id": "2"}`))

	// This should be dropped
	agent.QueueEvent("graphql://subscription/taskCreated", json.RawMessage(`{"id": "3"}`))

	stats := agent.GetQueueStats()
	if stats.TaskQueueSize != 2 {
		t.Errorf("Expected task queue size 2 (full), got %d", stats.TaskQueueSize)
	}
}

// TestAgent_StartStop tests the queue processor lifecycle
func TestAgent_StartStop(t *testing.T) {
	mockMCP := NewMockMCPClient([]Tool{
		{Name: "query", Description: "Execute a GraphQL query", InputSchema: InputSchema{Type: "object"}},
	})
	mockMCP.SetToolResult("query", &CallToolResult{
		Content: []ContentBlock{{Type: "text", Text: `{"data": {}}`}},
	})

	// Mock LLM that responds immediately
	mockLLM := NewMockLLM([]*openai.ChatCompletionMessage{
		{
			Role:    openai.ChatMessageRoleAssistant,
			Content: "Processed",
		},
	})

	agent := NewTestAgent(mockLLM, mockMCP)
	agent.Initialize()
	agent.SetupQueues(10)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Start the queue processor
	agent.Start(ctx)

	// Queue an event
	agent.QueueEvent("graphql://subscription/taskCreated", json.RawMessage(`{"id": "1"}`))

	// Give it time to process
	time.Sleep(100 * time.Millisecond)

	// Stop the processor
	agent.Stop()

	// Verify the queue is empty (event was processed)
	stats := agent.GetQueueStats()
	if stats.TaskQueueSize != 0 {
		t.Errorf("Expected task queue to be empty after processing, got %d", stats.TaskQueueSize)
	}
}

// TestAgent_MultipleEventsInOrder tests that events are processed in arrival order
func TestAgent_MultipleEventsInOrder(t *testing.T) {
	mockMCP := NewMockMCPClient([]Tool{
		{Name: "query", Description: "Execute a GraphQL query", InputSchema: InputSchema{Type: "object"}},
	})
	mockMCP.SetToolResult("query", &CallToolResult{
		Content: []ContentBlock{{Type: "text", Text: `{"data": {}}`}},
	})

	// Track the order of processed events
	var processedOrder []string
	// Mock LLM that responds immediately and tracks order
	mockLLM := NewMockLLM([]*openai.ChatCompletionMessage{
		{
			Role:    openai.ChatMessageRoleAssistant,
			Content: "Processed",
		},
		{
			Role:    openai.ChatMessageRoleAssistant,
			Content: "Processed",
		},
		{
			Role:    openai.ChatMessageRoleAssistant,
			Content: "Processed",
		},
	})

	agent := NewTestAgent(mockLLM, mockMCP)
	agent.Initialize()
	agent.SetupQueues(10)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Start the queue processor
	agent.Start(ctx)

	// Queue multiple events in order: task, message, task
	agent.QueueEvent("graphql://subscription/taskCreated", json.RawMessage(`{"id": "task-1"}`))
	agent.QueueEvent("graphql://subscription/messageAdded", json.RawMessage(`{"id": "msg-1"}`))
	agent.QueueEvent("graphql://subscription/taskUpdated", json.RawMessage(`{"id": "task-2"}`))

	// Give it time to process all events
	time.Sleep(200 * time.Millisecond)

	// Stop the processor
	agent.Stop()

	// Verify all queues are empty
	stats := agent.GetQueueStats()
	if stats.TaskQueueSize != 0 {
		t.Errorf("Expected task queue to be empty, got %d", stats.TaskQueueSize)
	}
	if stats.MessageQueueSize != 0 {
		t.Errorf("Expected message queue to be empty, got %d", stats.MessageQueueSize)
	}

	// Verify order was preserved (we can't easily check this with the mock, but the test validates the mechanism)
	_ = processedOrder
}