package gateway import ( "context" "io" "log/slog" "net/http" "net/http/httptest" "os" "strings" "sync/atomic" "testing" "time" "aigateway.local/core/internal/audit" "aigateway.local/core/internal/contentpolicy" "aigateway.local/core/internal/pricing" provideropenai "aigateway.local/core/internal/provider/openai" "github.com/jackc/pgx/v5/pgxpool" ) func TestContentPolicyAndPricingIntegration(t *testing.T) { databaseURL := os.Getenv("CONTENT_PRICING_TEST_DATABASE_URL") if databaseURL == "" { t.Skip("CONTENT_PRICING_TEST_DATABASE_URL is not set") } ctx := context.Background() pool, err := pgxpool.New(ctx, databaseURL) if err != nil { t.Fatal(err) } defer pool.Close() const blockID = "00000000-0000-4000-8000-000000001601" const priceID = "00000000-0000-4000-8000-000000001602" _, err = pool.Exec(ctx, `INSERT INTO gateway.content_policies(id,name,action,priority,rules,enabled) VALUES($1,'integration block','block',2000,'[{"name":"forbidden","pattern":"forbidden","replacement":"[BLOCKED]"}]',true) ON CONFLICT(id) DO UPDATE SET enabled=true`, blockID) if err != nil { t.Fatal(err) } _, err = pool.Exec(ctx, `INSERT INTO gateway.model_prices(id,provider_code,model_pattern,input_microunits_per_million,output_microunits_per_million,currency,effective_from,enabled) VALUES($1,'environment','gpt-test',2000000,8000000,'USD',clock_timestamp()-interval '1 hour',true) ON CONFLICT(id) DO UPDATE SET enabled=true`, priceID) if err != nil { t.Fatal(err) } defer func() { _, _ = pool.Exec(context.Background(), `DELETE FROM gateway.outbox_events WHERE event_type='content_policy.matched' AND aggregate_id IN ('m3-content-price-redact','m3-content-price-block')`) _, _ = pool.Exec(context.Background(), `DELETE FROM gateway.content_policies WHERE id=$1`, blockID) _, _ = pool.Exec(context.Background(), `DELETE FROM gateway.model_prices WHERE id=$1`, priceID) }() var upstreamCalls atomic.Int64 bodySeen := make(chan string, 1) upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { upstreamCalls.Add(1) body, _ := io.ReadAll(r.Body) bodySeen <- string(body) w.Header().Set("Content-Type", "application/json") _, _ = io.WriteString(w, `{"id":"ok","usage":{"prompt_tokens":1000,"completion_tokens":500,"total_tokens":1500}}`) })) defer upstream.Close() adapter, _ := provideropenai.New(upstream.URL, "upstream-key") engine := contentpolicy.NewEngine(pool, time.Minute, slog.Default()) if err := engine.Reload(ctx); err != nil { t.Fatal(err) } prices := pricing.NewService(pool, time.Minute, slog.Default()) if err := prices.Reload(ctx); err != nil { t.Fatal(err) } recorder := audit.NewRecorder(pool, slog.Default(), 100, 1, 10*time.Millisecond) recordCtx, cancel := context.WithCancel(ctx) stopped := make(chan struct{}) go func() { recorder.Run(recordCtx); close(stopped) }() defer func() { cancel(); <-stopped }() proxy := NewProxy(adapter, "test-key", 1<<20, slog.Default()) proxy.SetAllowPrivateProviderURLs(true) // 测试上游监听 127.0.0.1 proxy.SetAuditRecorder(recorder) proxy.SetContentPolicyEngine(engine) proxy.SetPricingService(prices) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { proxy.ServeHTTP(w, r.WithContext(WithRequestID(r.Context(), r.Header.Get("X-Request-ID")))) })) defer server.Close() req, _ := http.NewRequest(http.MethodPost, server.URL+"/v1/chat/completions", strings.NewReader(`{"model":"gpt-test","messages":[{"role":"user","content":"token sk-1234567890123456"}]}`)) req.Header.Set("Authorization", "Bearer test-key") req.Header.Set("X-Request-ID", "m3-content-price-redact") response, err := http.DefaultClient.Do(req) if err != nil { t.Fatal(err) } _, _ = io.ReadAll(response.Body) _ = response.Body.Close() if response.StatusCode != 200 || response.Header.Get("X-Gateway-Content-Redacted") != "true" { t.Fatalf("unexpected response %d headers=%v", response.StatusCode, response.Header) } select { case body := <-bodySeen: if strings.Contains(body, "sk-1234567890123456") || !strings.Contains(body, "[REDACTED_API_KEY]") { t.Fatalf("upstream saw unsafe body %s", body) } case <-time.After(time.Second): t.Fatal("upstream did not receive request") } blocked, _ := http.NewRequest(http.MethodPost, server.URL+"/v1/responses", strings.NewReader(`{"model":"gpt-test","input":"forbidden"}`)) blocked.Header.Set("Authorization", "Bearer test-key") blocked.Header.Set("X-Request-ID", "m3-content-price-block") blockedResponse, err := http.DefaultClient.Do(blocked) if err != nil { t.Fatal(err) } _, _ = io.ReadAll(blockedResponse.Body) _ = blockedResponse.Body.Close() if blockedResponse.StatusCode != http.StatusUnprocessableEntity || upstreamCalls.Load() != 1 { t.Fatalf("block failed status=%d calls=%d", blockedResponse.StatusCode, upstreamCalls.Load()) } deadline := time.Now().Add(2 * time.Second) for { var cost *int64 var labels string err = pool.QueryRow(ctx, `SELECT cost_microunits,labels::text FROM gateway.audit_events WHERE request_id='m3-content-price-redact' ORDER BY recorded_at DESC LIMIT 1`).Scan(&cost, &labels) if err == nil && cost != nil && *cost == 6000 && strings.Contains(labels, "content_redacted") { break } if time.Now().After(deadline) { t.Fatalf("audit cost/redaction missing cost=%v labels=%s err=%v", cost, labels, err) } time.Sleep(20 * time.Millisecond) } deadline = time.Now().Add(2 * time.Second) for { var count int err = pool.QueryRow(ctx, `SELECT count(*) FROM gateway.outbox_events WHERE event_type='content_policy.matched' AND aggregate_id IN ('m3-content-price-redact','m3-content-price-block')`).Scan(&count) if err == nil && count == 2 { break } if time.Now().After(deadline) { t.Fatalf("content policy notification outbox missing count=%d err=%v", count, err) } time.Sleep(20 * time.Millisecond) } }