rony-llm-agent/pkg/llm/mock/mock_test.go

package mock

import (
	"context"
	"fmt"
	"testing"

	"github.com/VictorVargas/rony-llm-agent/pkg/llm"
)

func TestMockLLMClient_Generate_Default(t *testing.T) {
	client := New()
	resp, err := client.Generate(context.Background(), llm.CompletionRequest{})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if resp.Content != "default response" {
		t.Errorf("expected 'default response', got %q", resp.Content)
	}
	if resp.StopReason != llm.StopReasonEndTurn {
		t.Errorf("expected StopReasonEndTurn, got %q", resp.StopReason)
	}
}

func TestMockLLMClient_Generate_Custom(t *testing.T) {
	expected := llm.CompletionResponse{
		ID:         "test_123",
		Model:      "gpt-4",
		Content:    "hello world",
		StopReason: llm.StopReasonToolUse,
	}
	client := NewWithGenerate(expected)

	resp, err := client.Generate(context.Background(), llm.CompletionRequest{})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if resp.ID != "test_123" {
		t.Errorf("expected ID 'test_123', got %q", resp.ID)
	}
	if resp.Content != "hello world" {
		t.Errorf("expected 'hello world', got %q", resp.Content)
	}
	if resp.StopReason != llm.StopReasonToolUse {
		t.Errorf("expected StopReasonToolUse, got %q", resp.StopReason)
	}
}

func TestMockLLMClient_Generate_Error(t *testing.T) {
	client := &MockLLMClient{
		GenerateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
			return llm.CompletionResponse{}, fmt.Errorf("test error")
		},
	}

	_, err := client.Generate(context.Background(), llm.CompletionRequest{})
	if err == nil {
		t.Fatal("expected error, got nil")
	}
	if err.Error() != "test error" {
		t.Errorf("expected 'test error', got %q", err.Error())
	}
}

func TestMockLLMClient_Stream_Default(t *testing.T) {
	client := New()
	var chunks []string
	for chunk, err := range client.Stream(context.Background(), llm.CompletionRequest{}) {
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		chunks = append(chunks, chunk.Delta)
	}
	if len(chunks) != 1 || chunks[0] != "default response" {
		t.Errorf("expected single 'default response' chunk, got %v", chunks)
	}
}

func TestMockLLMClient_Stream_Custom(t *testing.T) {
	expectedChunks := []llm.StreamChunk{
		{Delta: "chunk1"},
		{Delta: "chunk2", ToolCalls: []llm.ToolCall{{ID: "call_1"}}},
		{FinishReason: "stop", Usage: llm.TokenUsage{InputTokens: 5, OutputTokens: 10}},
	}
	client := NewWithStream(expectedChunks)

	var actualChunks []llm.StreamChunk
	for chunk, err := range client.Stream(context.Background(), llm.CompletionRequest{}) {
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		actualChunks = append(actualChunks, chunk)
	}

	if len(actualChunks) != 3 {
		t.Errorf("expected 3 chunks, got %d", len(actualChunks))
		return
	}
	if actualChunks[0].Delta != "chunk1" {
		t.Errorf("expected 'chunk1', got %q", actualChunks[0].Delta)
	}
	if len(actualChunks[1].ToolCalls) != 1 {
		t.Errorf("expected 1 tool call, got %d", len(actualChunks[1].ToolCalls))
	}
	if actualChunks[2].FinishReason != "stop" {
		t.Errorf("expected 'stop', got %q", actualChunks[2].FinishReason)
	}
}

func TestMockLLMClient_Name_Default(t *testing.T) {
	client := New()
	if name := client.Name(); name != "mock" {
		t.Errorf("expected 'mock', got %q", name)
	}
}

func TestMockLLMClient_Name_Custom(t *testing.T) {
	client := &MockLLMClient{
		NameFunc: func() string { return "custom-model" },
	}
	if name := client.Name(); name != "custom-model" {
		t.Errorf("expected 'custom-model', got %q", name)
	}
}

func TestMockLLMClient_Capabilities_Default(t *testing.T) {
	client := New()
	caps := client.Capabilities()
	if !caps.SupportsTools {
		t.Error("expected SupportsTools to be true")
	}
	if caps.SupportsVision {
		t.Error("expected SupportsVision to be false")
	}
	if !caps.SupportsJSON {
		t.Error("expected SupportsJSON to be true")
	}
	if caps.MaxContextWindow != 128000 {
		t.Errorf("expected MaxContextWindow 128000, got %d", caps.MaxContextWindow)
	}
}

func TestMockLLMClient_Capabilities_Custom(t *testing.T) {
	client := &MockLLMClient{
		CapabilitiesFunc: func() llm.ProviderCapabilities {
			return llm.ProviderCapabilities{
				SupportsTools:    false,
				SupportsVision:   true,
				SupportsJSON:     false,
				MaxContextWindow: 32000,
			}
		},
	}
	caps := client.Capabilities()
	if caps.SupportsTools {
		t.Error("expected SupportsTools to be false")
	}
	if !caps.SupportsVision {
		t.Error("expected SupportsVision to be true")
	}
	if caps.MaxContextWindow != 32000 {
		t.Errorf("expected MaxContextWindow 32000, got %d", caps.MaxContextWindow)
	}
}

func TestMockLLMClient_MatchResponse(t *testing.T) {
	client := NewWithMatch([]MatchResponse{
		{Match: "hello", Response: "hi there!"},
		{Match: "world", Response: "earth"},
		{Match: "*",    Response: "default"},
	})

	resp, _ := client.Generate(context.Background(), llm.CompletionRequest{
		Messages: []llm.Message{{Role: llm.RoleUser, Content: "hello"}},
	})
	if resp.Content != "hi there!" {
		t.Errorf("expected 'hi there!', got %q", resp.Content)
	}

	resp, _ = client.Generate(context.Background(), llm.CompletionRequest{
		Messages: []llm.Message{{Role: llm.RoleUser, Content: "world"}},
	})
	if resp.Content != "earth" {
		t.Errorf("expected 'earth', got %q", resp.Content)
	}

	resp, _ = client.Generate(context.Background(), llm.CompletionRequest{
		Messages: []llm.Message{{Role: llm.RoleUser, Content: "anything else"}},
	})
	if resp.Content != "default" {
		t.Errorf("expected 'default', got %q", resp.Content)
	}
}

func TestMockLLMClient_MatchResponse_NoMatch(t *testing.T) {
	client := NewWithMatch([]MatchResponse{
		{Match: "exact", Response: "matched"},
	})

	resp, _ := client.Generate(context.Background(), llm.CompletionRequest{
		Messages: []llm.Message{{Role: llm.RoleUser, Content: "no match here"}},
	})
	if resp.Content != "no match" {
		t.Errorf("expected 'no match', got %q", resp.Content)
	}
}

func TestMatchesAny(t *testing.T) {
	if !matchesAny("hello", []llm.Message{
		{Role: llm.RoleUser, Content: "say hello to me"},
	}) {
		t.Error("expected matchesAny to find 'hello' in message")
	}

	if matchesAny("goodbye", []llm.Message{
		{Role: llm.RoleUser, Content: "say hello to me"},
	}) {
		t.Error("expected matchesAny to not find 'goodbye' in message")
	}
}

func TestMatchesAny_Wildcard(t *testing.T) {
	resp, _ := (&MockLLMClient{
		GenerateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
			for _, r := range []MatchResponse{{Match: "*", Response: "wildcard"}} {
				if r.Match == "*" || matchesAny(r.Match, req.Messages) {
					return llm.CompletionResponse{Content: r.Response}, nil
				}
			}
			return llm.CompletionResponse{}, nil
		},
	}).Generate(context.Background(), llm.CompletionRequest{})

	if resp.Content != "wildcard" {
		t.Errorf("expected 'wildcard', got %q", resp.Content)
	}
}