rony-llm-agent/pkg/agent/integration_test.go

package agent_test

import (
	"context"
	"encoding/json"
	"os"
	"testing"

	"github.com/VictorVargas/rony-llm-agent/pkg/agent"
	"github.com/VictorVargas/rony-llm-agent/pkg/llm"
	"github.com/VictorVargas/rony-llm-agent/pkg/llm/providers/llamacpp"
	"github.com/VictorVargas/rony-llm-agent/pkg/persona"
	"github.com/VictorVargas/rony-llm-agent/pkg/tools"
)

// TestIntegration_LlamaCPP_Generate is an integration test that requires llama.cpp running on localhost:8080.
// Run with: INTEGRATION_TESTS=1 go test ./pkg/agent/ -run TestIntegration_LlamaCPP_Generate
func TestIntegration_LlamaCPP_Generate(t *testing.T) {
	skipUnlessIntegration(t)
	client, err := llamacpp.New(llamacpp.Config{
		BaseURL: "http://localhost:8080/v1",
	})
	if err != nil {
		t.Fatalf("failed to create client: %v", err)
	}

	resp, err := client.Generate(context.Background(), llm.CompletionRequest{
		Messages: []llm.Message{
			{Role: llm.RoleUser, Content: "What is 2+2? Answer with just the number."},
		},
	})
	if err != nil {
		t.Fatalf("generate failed: %v", err)
	}

	if resp.Content == "" {
		t.Fatal("expected non-empty response")
	}
}

// TestIntegration_LlamaCPP_Stream is an integration test that requires llama.cpp running on localhost:8080.
func TestIntegration_LlamaCPP_Stream(t *testing.T) {
	skipUnlessIntegration(t)
	client, err := llamacpp.New(llamacpp.Config{
		BaseURL: "http://localhost:8080/v1",
	})
	if err != nil {
		t.Fatalf("failed to create client: %v", err)
	}

	stream := client.Stream(context.Background(), llm.CompletionRequest{
		Messages: []llm.Message{
			{Role: llm.RoleUser, Content: "Say hello in 5 words."},
		},
	})

	var chunks []llm.StreamChunk
	for chunk, err := range stream {
		if err != nil {
			t.Fatalf("stream error: %v", err)
		}
		chunks = append(chunks, chunk)
	}

	if len(chunks) == 0 {
		t.Fatal("expected at least one chunk")
	}
}

// TestIntegration_AgentLoop_Generate is an integration test for the agent loop with llama.cpp.
func TestIntegration_AgentLoop_Generate(t *testing.T) {
	skipUnlessIntegration(t)
	client, err := llamacpp.New(llamacpp.Config{
		BaseURL: "http://localhost:8080/v1",
	})
	if err != nil {
		t.Fatalf("failed to create client: %v", err)
	}

	registry := tools.NewRegistry()
	registry.Register(tools.Tool{
		Name:        "add_numbers",
		Description: "Add two numbers together",
		InputSchema: json.RawMessage(`{"type": "function", "function": {"name": "add_numbers", "description": "Add two numbers together", "parameters": {"type": "object", "properties": {"a": {"type": "number"}, "b": {"type": "number"}}}}}`),
		Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) {
			var input struct {
				A float64 `json:"a"`
				B float64 `json:"b"`
			}
			if err := json.Unmarshal(args, &input); err != nil {
				return tools.ToolResult{IsError: true}, err
			}
			return tools.ToolResult{Content: "42"}, nil
		},
		Permission: tools.Allow,
	})

	loop := agent.New(agent.Config{
		LLM:      client,
		Persona:  persona.DefaultPersona(),
		Tools:    registry,
		Sandbox:  &mockSandbox{},
		MaxIters: 3,
	})

	resp, err := loop.Run(context.Background(), "What is 20+22? Use the add_numbers tool.")
	if err != nil {
		t.Fatalf("run failed: %v", err)
	}

	if resp.Content == "" {
		t.Fatal("expected non-empty response")
	}
}

// TestIntegration_AgentLoop_Stream is an integration test for the agent loop with streaming.
func TestIntegration_AgentLoop_Stream(t *testing.T) {
	skipUnlessIntegration(t)
	client, err := llamacpp.New(llamacpp.Config{
		BaseURL: "http://localhost:8080/v1",
	})
	if err != nil {
		t.Fatalf("failed to create client: %v", err)
	}

	loop := agent.New(agent.Config{
		LLM:      client,
		Persona:  persona.DefaultPersona(),
		Tools:    tools.NewRegistry(),
		MaxIters: 3,
	})

	stream := loop.RunStream(context.Background(), "Say something interesting.")

	var chunks []llm.StreamChunk
	for chunk, err := range stream {
		if err != nil {
			t.Fatalf("stream error: %v", err)
		}
		chunks = append(chunks, chunk)
	}

	if len(chunks) == 0 {
		t.Fatal("expected at least one chunk")
	}
}

// mockSandbox is a simple sandbox that allows all calls.
type mockSandbox struct{}

func (m *mockSandbox) ValidateToolCall(tool tools.Tool, call llm.ToolCall) error {
	return nil
}

// skipUnlessIntegration skips a single integration test (one that needs a
// live llama.cpp server on localhost:8080) unless explicitly enabled.
//
// This used to be done in TestMain by returning early without calling
// m.Run() when INTEGRATION_TESTS wasn't set - but a package's TestMain
// covers its *entire* test binary (both package agent_test, here, and
// package agent, e.g. loop_test.go, get linked together), so that skipped
// every test in the package, not just the four integration ones. In
// practice that meant `go test ./...` reported this package as passing
// while silently running zero of its tests, including all the mock-based
// coverage in loop_test.go for approvals, sandboxing, and tool-call
// handling.
func skipUnlessIntegration(t *testing.T) {
	t.Helper()
	if os.Getenv("INTEGRATION_TESTS") != "1" {
		t.Skip("set INTEGRATION_TESTS=1 to run (requires a live llama.cpp server on localhost:8080)")
	}
}