rsspipe/internal/ai/client_test.go

package ai

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"testing"
	"time"

	"git.pengzhan.dev/haopengzhan/rsspipe/internal/models"
)

type roundTripFunc func(req *http.Request) *http.Response

func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
	return f(req), nil
}

func TestTestConnection(t *testing.T) {
	client := &Client{
		httpClient: &http.Client{
			Transport: roundTripFunc(func(req *http.Request) *http.Response {
				if req.Header.Get("Authorization") != "Bearer test-key" {
					return &http.Response{
						StatusCode: http.StatusUnauthorized,
						Body:       io.NopCloser(bytes.NewBufferString(`{"error": {"message": "unauthorized"}}`)),
						Header:     make(http.Header),
					}
				}

				resp := chatCompletionResponse{
					Choices: []struct {
						Message struct {
							Role    string `json:"role"`
							Content string `json:"content"`
						} `json:"message"`
						FinishReason string `json:"finish_reason"`
					}{
						{
							Message: struct {
								Role    string `json:"role"`
								Content string `json:"content"`
							}{
								Role:    "assistant",
								Content: "OK connection",
							},
							FinishReason: "stop",
						},
					},
				}
				data, _ := json.Marshal(resp)
				return &http.Response{
					StatusCode: http.StatusOK,
					Body:       io.NopCloser(bytes.NewReader(data)),
					Header:     make(http.Header),
				}
			}),
		},
	}

	cfg := Config{
		Enabled:  true,
		Provider: "custom",
		BaseURL:  "https://mock.api",
		APIKey:   "test-key",
		Model:    "test-model",
	}

	res, err := client.TestConnection(context.Background(), cfg)
	if err != nil {
		t.Fatalf("TestConnection failed: %v", err)
	}
	if !res.Success {
		t.Fatalf("Expected success true, got false")
	}
	if res.Message != "OK connection" {
		t.Fatalf("Expected 'OK connection', got %s", res.Message)
	}
}

func TestGenerateEnrichment(t *testing.T) {
	client := &Client{
		httpClient: &http.Client{
			Transport: roundTripFunc(func(req *http.Request) *http.Response {
				resp := chatCompletionResponse{
					Choices: []struct {
						Message struct {
							Role    string `json:"role"`
							Content string `json:"content"`
						} `json:"message"`
						FinishReason string `json:"finish_reason"`
					}{
						{
							Message: struct {
								Role    string `json:"role"`
								Content string `json:"content"`
							}{
								Role:    "assistant",
								Content: "- Point 1\n- Point 2\n- Point 3\n\n**Key Takeaway**: Great post.",
							},
							FinishReason: "stop",
						},
					},
				}
				data, _ := json.Marshal(resp)
				return &http.Response{
					StatusCode: http.StatusOK,
					Body:       io.NopCloser(bytes.NewReader(data)),
					Header:     make(http.Header),
				}
			}),
		},
	}

	cfg := Config{
		Enabled:   true,
		BaseURL:   "https://mock.api",
		Model:     "gpt-4o-mini",
		MaxTokens: 500,
	}

	entry := models.FeedEntry{
		ID:          "e-1",
		Title:       "Kubernetes & GPU Optimization",
		FeedTitle:   "Cloud Native Blog",
		Content:     "Detailed technical content about GPU sharing in Kubernetes.",
		PublishedAt: time.Now(),
	}

	summary, err := client.GenerateEnrichment(context.Background(), cfg, "Summarize: {{title}}\n{{content}}", entry)
	if err != nil {
		t.Fatalf("GenerateEnrichment failed: %v", err)
	}
	if summary == "" {
		t.Fatalf("Expected summary string, got empty")
	}
}

func TestGenerateAgentDecision(t *testing.T) {
	client := &Client{
		httpClient: &http.Client{
			Transport: roundTripFunc(func(req *http.Request) *http.Response {
				resp := chatCompletionResponse{
					Choices: []struct {
						Message struct {
							Role    string `json:"role"`
							Content string `json:"content"`
						} `json:"message"`
						FinishReason string `json:"finish_reason"`
					}{
						{
							Message: struct {
								Role    string `json:"role"`
								Content string `json:"content"`
							}{
								Role: "assistant",
								Content: `{
									"reasoning": "High priority kubernetes security patch",
									"actions": [
										{"type": "mark_starred"},
										{"type": "set_variant", "params": {"variant_key": "tags", "value": "security-alert"}},
										{"type": "malicious_unallowed_action"}
									]
								}`,
							},
							FinishReason: "stop",
						},
					},
				}
				data, _ := json.Marshal(resp)
				return &http.Response{
					StatusCode: http.StatusOK,
					Body:       io.NopCloser(bytes.NewReader(data)),
					Header:     make(http.Header),
				}
			}),
		},
	}

	cfg := Config{
		Enabled: true,
		BaseURL: "https://mock.api",
		Model:   "gpt-4o-mini",
	}

	entry := models.FeedEntry{
		ID:        "e-1",
		Title:     "Critical CVE in Kubelet",
		FeedTitle: "Kubernetes Security",
		Content:   "Remote code execution flaw patched in v1.36.",
	}

	allowed := []string{"mark_starred", "set_variant:tags:security-alert"}
	decision, err := client.GenerateAgentDecision(context.Background(), cfg, "Evaluate urgency", allowed, entry)
	if err != nil {
		t.Fatalf("GenerateAgentDecision failed: %v", err)
	}

	if decision.Reasoning == "" {
		t.Fatalf("Expected reasoning, got empty")
	}

	// Verify strict validation: malicious_unallowed_action was stripped
	if len(decision.Actions) != 2 {
		t.Fatalf("Expected 2 validated actions, got %d", len(decision.Actions))
	}
	if decision.Actions[0].Type != "mark_starred" {
		t.Fatalf("Expected mark_starred, got %s", decision.Actions[0].Type)
	}
}