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) } }