ai-gateway-go/internal/gateway/content_pricing_integration_test.go

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