zarinpal-analytics-backend/internal/growthreport/avalai_test.go

package growthreport

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func TestAvalAICountsFailedRetryRequestsAndUsage(t *testing.T) {
	t.Parallel()
	requests := 0
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		requests++
		w.Header().Set("X-Request-ID", fmt.Sprintf("request-%d", requests))
		if requests == 1 {
			w.WriteHeader(http.StatusServiceUnavailable)
			_, _ = w.Write([]byte(`{"usage":{"input_tokens":7,"output_tokens":3},"estimated_cost":"0.01"}`))
			return
		}
		_, _ = w.Write([]byte(`{"status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"{\"briefs\":[{\"headline\":\"فرصت مناسب\",\"interpretation\":\"هم‌راستا با شواهد است.\",\"offerSuggestion\":\"پیشنهاد آماده کنید.\",\"messageSuggestion\":\"پیام روشن ارسال کنید.\",\"preparationAction\":\"مسیر پرداخت را بررسی کنید.\",\"limitation\":\"نتیجه قطعی نیست.\",\"evidenceIds\":[\"event_teacher\"]}]}"}]}],"usage":{"input_tokens":11,"output_tokens":5},"estimated_cost":"0.02"}`))
	}))
	defer server.Close()

	client := NewAvalAIClient(AvalAIConfig{BaseURL: server.URL, APIKey: "test", Model: "test", Timeout: time.Second}, server.Client())
	_, usage, err := client.Generate(context.Background(), json.RawMessage(`{"evidence":"event_teacher"}`), map[string]struct{}{"event_teacher": {}})
	if err != nil {
		t.Fatal(err)
	}
	if requests != 2 || usage.ProviderRequests != 2 || usage.InputTokens != 18 || usage.OutputTokens != 8 || usage.EstimatedCost != "0.03" {
		t.Fatalf("requests=%d usage=%+v", requests, usage)
	}
}

func TestValidateAICopyRejectsNumbersAndUnknownEvidence(t *testing.T) {
	t.Parallel()
	tests := []struct {
		name string
		copy AICopy
		want bool
	}{
		{name: "grounded Persian copy", copy: AICopy{Briefs: []CampaignBrief{{Headline: "فرصت مناسبی پیش رو دارید", Interpretation: "این مناسبت با الگوی موفق گذشته هم‌راستا است.", EvidenceIDs: []string{"event_teacher"}}}}, want: true},
		{name: "numeric claim", copy: AICopy{Briefs: []CampaignBrief{{Headline: "رشد ۲۰ درصدی", EvidenceIDs: []string{"event_teacher"}}}}, want: false},
		{name: "Persian number-word claim", copy: AICopy{Briefs: []CampaignBrief{{Headline: "رشد بیست درصدی", EvidenceIDs: []string{"event_teacher"}}}}, want: false},
		{name: "English number-word claim", copy: AICopy{Briefs: []CampaignBrief{{Headline: "growth by twenty percent", EvidenceIDs: []string{"event_teacher"}}}}, want: false},
		{name: "unknown evidence", copy: AICopy{Briefs: []CampaignBrief{{Headline: "فرصت مناسب", EvidenceIDs: []string{"unknown"}}}}, want: false},
	}
	known := map[string]struct{}{"event_teacher": {}}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			err := ValidateAICopy(tt.copy, known)
			if (err == nil) != tt.want {
				t.Fatalf("ValidateAICopy() error = %v, valid want %v", err, tt.want)
			}
		})
	}
}

func TestAvalAIClientUsesBoundedStatelessRequest(t *testing.T) {
	t.Parallel()
	var request map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer test-key" {
			t.Errorf("authorization header missing")
		}
		if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
			t.Fatal(err)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"{\"briefs\":[{\"headline\":\"فرصت مناسبی پیش رو دارید\",\"interpretation\":\"با الگوی موفق گذشته هم‌راستا است.\",\"offerSuggestion\":\"پیشنهاد مرتبط آماده کنید.\",\"messageSuggestion\":\"پیام کوتاه و روشن ارسال کنید.\",\"preparationAction\":\"ظرفیت پرداخت را بررسی کنید.\",\"limitation\":\"این پیشنهاد قطعی نیست.\",\"evidenceIds\":[\"event_teacher\"]}]}"}]}],"usage":{"input_tokens":10,"output_tokens":20}}`))
	}))
	defer server.Close()

	client := NewAvalAIClient(AvalAIConfig{BaseURL: server.URL, APIKey: "test-key", Model: "test-model", Timeout: time.Second}, server.Client())
	copy, usage, err := client.Generate(context.Background(), json.RawMessage(`{"evidenceIds":["event_teacher"]}`), map[string]struct{}{"event_teacher": {}})
	if err != nil {
		t.Fatal(err)
	}
	if len(copy.Briefs) != 1 || usage.InputTokens != 10 || usage.OutputTokens != 20 {
		t.Fatalf("Generate() = %+v, %+v", copy, usage)
	}
	if request["store"] != false || request["parallel_tool_calls"] != false {
		t.Fatalf("unsafe request controls: %#v", request)
	}
	if request["max_output_tokens"] != float64(1600) {
		t.Fatalf("max_output_tokens = %#v", request["max_output_tokens"])
	}
	tools, ok := request["tools"].([]any)
	if !ok || len(tools) != 4 {
		t.Fatalf("tools = %#v", request["tools"])
	}
	encoded, _ := json.Marshal(request)
	if strings.Contains(string(encoded), "merchantId") || strings.Contains(string(encoded), "payer") {
		t.Fatalf("sensitive field in request: %s", encoded)
	}
}

type testTools struct{ calls int }

func (t *testTools) Execute(_ context.Context, name string, _ json.RawMessage) (json.RawMessage, error) {
	t.calls++
	return json.RawMessage(`{"state":"READY"}`), nil
}

func TestAvalAIClientExecutesAllowlistedToolThenStopsAfterSecondTurn(t *testing.T) {
	t.Parallel()
	turn := 0
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		turn++
		w.Header().Set("Content-Type", "application/json")
		if turn == 1 {
			_, _ = w.Write([]byte(`{"status":"completed","output":[{"type":"function_call","call_id":"call_1","name":"get_growth_context","arguments":"{\"evidenceId\":\"event_teacher\"}"}],"usage":{"input_tokens":2,"output_tokens":1}}`))
			return
		}
		_, _ = w.Write([]byte(`{"status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"{\"briefs\":[{\"headline\":\"فرصت مناسب\",\"interpretation\":\"هم‌راستا با شواهد است.\",\"offerSuggestion\":\"پیشنهاد مرتبط آماده کنید.\",\"messageSuggestion\":\"پیام روشن ارسال کنید.\",\"preparationAction\":\"مسیر پرداخت را بررسی کنید.\",\"limitation\":\"نتیجه قطعی نیست.\",\"evidenceIds\":[\"event_teacher\"]}]}"}]}],"usage":{"input_tokens":3,"output_tokens":4}}`))
	}))
	defer server.Close()
	executor := &testTools{}
	client := NewAvalAIClient(AvalAIConfig{BaseURL: server.URL, APIKey: "key", Model: "model", Timeout: time.Second}, server.Client())
	_, usage, err := client.GenerateWithTools(context.Background(), json.RawMessage(`{}`), map[string]struct{}{"event_teacher": {}}, executor)
	if err != nil {
		t.Fatal(err)
	}
	if executor.calls != 1 || turn != 2 || usage.InputTokens != 5 || usage.OutputTokens != 5 {
		t.Fatalf("calls=%d turns=%d usage=%+v", executor.calls, turn, usage)
	}
}