rony-llm-agent/pkg/llm/providers/openai/client_test.go

package openai

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

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

func TestClient_Name(t *testing.T) {
	client := &Client{}
	if client.Name() != "openai" {
		t.Errorf("expected 'openai', got %q", client.Name())
	}
}

func TestClient_Capabilities(t *testing.T) {
	client := &Client{}
	caps := client.Capabilities()
	if !caps.SupportsTools {
		t.Error("expected SupportsTools to be true")
	}
	if !caps.SupportsVision {
		t.Error("expected SupportsVision to be true")
	}
}

func TestClient_Generate(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		json.NewEncoder(w).Encode(openaiChatResponse{
			ID:    "test-123",
			Model: "gpt-4",
			Choices: []openaiChoice{
				{
					Index:        0,
					FinishReason: "stop",
					Message: openaiMessageResult{
						Role:    "assistant",
						Content: "Hello!",
					},
				},
			},
			Usage: openaiUsage{
				PromptTokens:     10,
				CompletionTokens: 5,
				TotalTokens:      15,
			},
		})
	}))
	defer server.Close()

	client, err := New(Config{
		APIKey:  "test-key",
		BaseURL: server.URL,
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	resp, err := client.Generate(context.Background(), llm.CompletionRequest{
		Model: "gpt-4",
		Messages: []llm.Message{
			{Role: llm.RoleUser, Content: "Hi"},
		},
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if resp.Content != "Hello!" {
		t.Errorf("expected 'Hello!', got %q", resp.Content)
	}
	if resp.ID != "test-123" {
		t.Errorf("expected 'test-123', got %q", resp.ID)
	}
	if resp.Usage.InputTokens != 10 {
		t.Errorf("expected 10 input tokens, got %d", resp.Usage.InputTokens)
	}
}

func TestClient_Generate_Error(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusUnauthorized)
		w.Write([]byte("invalid api key"))
	}))
	defer server.Close()

	client, err := New(Config{
		APIKey:  "bad-key",
		BaseURL: server.URL,
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	_, err = client.Generate(context.Background(), llm.CompletionRequest{})
	if err == nil {
		t.Fatal("expected error, got nil")
	}
}

func TestClient_Stream(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n"))
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\" world\"},\"finish_reason\":null}]}\n"))
		w.Write([]byte("data: [DONE]\n"))
	}))
	defer server.Close()

	client, err := New(Config{
		APIKey:  "test-key",
		BaseURL: server.URL + "/v1",
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

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

	if len(chunks) != 2 {
		t.Errorf("expected 2 chunks, got %d", len(chunks))
	}
	if chunks[0].Delta != "Hello" {
		t.Errorf("expected 'Hello', got %q", chunks[0].Delta)
	}
	if chunks[1].Delta != " world" {
		t.Errorf("expected ' world', got %q", chunks[1].Delta)
	}
}