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":"{\"eventRecommendations\":[{\"evidenceId\":\"event_teacher\",\"relevanceReason\":\"این مناسبت با الگوی فروش کسب‌وکار هم‌خوان است.\"},{\"evidenceId\":\"event_yalda\",\"relevanceReason\":\"این رویداد با کانال فروش شما مرتبط است.\"},{\"evidenceId\":\"event_nowruz\",\"relevanceReason\":\"این بازه برای برنامه‌ریزی عرضه مناسب است.\"}]}"}]}],"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(`{"events":[]}`), map[string]struct{}{"event_teacher": {}, "event_yalda": {}, "event_nowruz": {}}) 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 TestValidateAICopyAcceptsExactlyThreeGroundedPersianRecommendations(t *testing.T) { t.Parallel() tests := []struct { name string copy AICopy wantErr bool }{ {name: "grounded Persian recommendations", copy: AICopy{EventRecommendations: []EventRecommendation{{EvidenceID: "event_teacher", RelevanceReason: "این مناسبت با الگوی فروش کسب‌وکار هم‌خوان است."}, {EvidenceID: "event_yalda", RelevanceReason: "این رویداد با کانال فروش شما مرتبط است."}, {EvidenceID: "event_nowruz", RelevanceReason: "این بازه برای برنامه‌ریزی عرضه مناسب است."}}}}, {name: "numeric claim", copy: AICopy{EventRecommendations: []EventRecommendation{{EvidenceID: "event_teacher", RelevanceReason: "رشد ۲۰ درصدی پیش‌بینی می‌شود."}, {EvidenceID: "event_yalda", RelevanceReason: "این رویداد با کانال فروش شما مرتبط است."}, {EvidenceID: "event_nowruz", RelevanceReason: "این بازه برای برنامه‌ریزی عرضه مناسب است."}}}, wantErr: true}, {name: "duplicate evidence", copy: AICopy{EventRecommendations: []EventRecommendation{{EvidenceID: "event_teacher", RelevanceReason: "این مناسبت با الگوی فروش کسب‌وکار هم‌خوان است."}, {EvidenceID: "event_teacher", RelevanceReason: "این رویداد با کانال فروش شما مرتبط است."}, {EvidenceID: "event_nowruz", RelevanceReason: "این بازه برای برنامه‌ریزی عرضه مناسب است."}}}, wantErr: true}, {name: "unknown evidence", copy: AICopy{EventRecommendations: []EventRecommendation{{EvidenceID: "event_teacher", RelevanceReason: "این مناسبت با الگوی فروش کسب‌وکار هم‌خوان است."}, {EvidenceID: "event_yalda", RelevanceReason: "این رویداد با کانال فروش شما مرتبط است."}, {EvidenceID: "unknown", RelevanceReason: "این بازه برای برنامه‌ریزی عرضه مناسب است."}}}, wantErr: true}, {name: "not exactly three", copy: AICopy{EventRecommendations: []EventRecommendation{{EvidenceID: "event_teacher", RelevanceReason: "این مناسبت با الگوی فروش کسب‌وکار هم‌خوان است."}}}, wantErr: true}, } known := map[string]struct{}{"event_teacher": {}, "event_yalda": {}, "event_nowruz": {}} for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { err := ValidateAICopy(tt.copy, known) if (err != nil) != tt.wantErr { t.Fatalf("ValidateAICopy() error = %v, wantErr %v", err, tt.wantErr) } }) } } 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":"{\"eventRecommendations\":[{\"evidenceId\":\"event_teacher\",\"relevanceReason\":\"این مناسبت با الگوی فروش کسب‌وکار هم‌خوان است.\"},{\"evidenceId\":\"event_yalda\",\"relevanceReason\":\"این رویداد با کانال فروش شما مرتبط است.\"},{\"evidenceId\":\"event_nowruz\",\"relevanceReason\":\"این بازه برای برنامه‌ریزی عرضه مناسب است.\"}]}"}]}],"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(`{"events":[{"evidenceId":"event_teacher"}]}`), map[string]struct{}{"event_teacher": {}, "event_yalda": {}, "event_nowruz": {}}) if err != nil { t.Fatal(err) } if len(copy.EventRecommendations) != 3 || 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":"{\"eventRecommendations\":[{\"evidenceId\":\"event_teacher\",\"relevanceReason\":\"این مناسبت با الگوی فروش کسب‌وکار هم‌خوان است.\"},{\"evidenceId\":\"event_yalda\",\"relevanceReason\":\"این رویداد با کانال فروش شما مرتبط است.\"},{\"evidenceId\":\"event_nowruz\",\"relevanceReason\":\"این بازه برای برنامه‌ریزی عرضه مناسب است.\"}]}"}]}],"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": {}, "event_yalda": {}, "event_nowruz": {}}, 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) } }