Files
new-api-custom/controller/detection_test.go
T

389 lines
19 KiB
Go

package controller
import (
"bufio"
"context"
"fmt"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
"gorm.io/driver/mysql"
"gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/schema"
"io"
"net/http"
"net/http/httptest"
"os"
"strings"
"testing"
"time"
)
func TestDetectionProxyFlushesBeforeGenerationAndConsumesCredentialOnce(t *testing.T) {
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
connection, err := db.DB()
require.NoError(t, err)
connection.SetMaxOpenConns(1)
oldDB := model.DB
model.DB = db
t.Cleanup(func() { model.DB = oldDB; _ = connection.Close() })
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}))
release := make(chan struct{})
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
body, readErr := io.ReadAll(r.Body)
require.NoError(t, readErr)
assert.Equal(t, "medium", gjson.GetBytes(body, "reasoning.effort").String())
select {
case <-release:
case <-r.Context().Done():
return
}
w.Header().Set("Content-Type", "application/json")
_, _ = io.WriteString(w, `{"status":"completed","usage":{"input_tokens":10,"output_tokens":5},"output":[{"type":"message","content":[{"type":"output_text","text":"<svg></svg>"}]}]}`)
}))
defer upstream.Close()
url := upstream.URL
channel := model.Channel{Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: fingerprint.Model}
require.NoError(t, db.Create(&channel).Error)
price, _ := common.Marshal(service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: "p+c", ExprHash: billingexpr.ExprHashString("p+c"), GroupRatio: 1, QuotaPerUnit: 500000, ExprVersion: 1}})
run := model.DetectionRun{ID: "stream-test", RequestKey: "stream-test", Kind: "pelican", ChannelID: channel.Id, Status: "running", Deadline: time.Now().Unix() + 60, Reserved: 100, Price: string(price)}
require.NoError(t, db.Create(&run).Error)
call := model.DetectionCall{ID: "stream-call", RunID: run.ID, SecretHash: model.DetectionSecretHash("det_test"), State: "ready", ExpiresAt: time.Now().Unix() + 60}
require.NoError(t, db.Create(&call).Error)
router := gin.New()
router.POST("/responses", DetectionProxy)
gateway := httptest.NewServer(router)
defer gateway.Close()
client := http.Client{Timeout: 5 * time.Second}
request, err := http.NewRequest(http.MethodPost, gateway.URL+"/responses", strings.NewReader(`{"model":"gpt-6-astra","input":"draw"}`))
require.NoError(t, err)
request.Header.Set("Authorization", "Bearer det_test")
response, err := client.Do(request)
close(release)
require.NoError(t, err)
defer response.Body.Close()
reader := bufio.NewReader(response.Body)
first, err := reader.ReadString('\n')
require.NoError(t, err)
assert.Equal(t, ": detection started\n", first)
raw, err := io.ReadAll(reader)
require.NoError(t, err)
assert.Contains(t, string(raw), "response.completed")
require.NoError(t, db.First(&call, "id = ?", call.ID).Error)
assert.Equal(t, "finished", call.State)
assert.Contains(t, call.Usage, "input_tokens")
replay := httptest.NewRecorder()
request = httptest.NewRequest(http.MethodPost, "/responses", strings.NewReader(`{"model":"gpt-6-astra","input":"draw"}`))
request.Header.Set("Authorization", "Bearer det_test")
router.ServeHTTP(replay, request)
assert.Equal(t, 401, replay.Code)
}
func TestDetectionProxyRejectsEscapesAndEnforcesBounds(t *testing.T) {
for _, body := range []string{
`{"model":[],"input":"hi"}`,
`{"model":"gpt-6-luna","input":"hi"}`,
`{"model":"gpt-6-astra","input":[{"role":[],"content":"hi"}]}`,
`{"model":"gpt-6-astra","input":"hi","previous_response_id":"x"}`,
`{"model":"gpt-6-astra","input":[{"type":"input_image","image_url":"file:///etc/passwd"}]}`,
} {
_, err := detectionProxyRequest([]byte(body))
require.Error(t, err)
}
request, err := detectionProxyRequest([]byte(`{"model":"gpt-6-astra","input":"draw a pelican","max_output_tokens":1000000,"tools":[{"type":"web_search"}]}`))
require.NoError(t, err)
require.NotNil(t, request.MaxOutputTokens)
assert.EqualValues(t, 16384, *request.MaxOutputTokens)
assert.Equal(t, "[]", string(request.Tools))
}
func TestDetectionOnlyBillsCompleteReportedUsage(t *testing.T) {
answer, parseErr := fingerprint.ReadResponse(strings.NewReader(`{"status":"incomplete","usage":{"input_tokens":100,"output_tokens":10,"input_tokens_details":{"cached_tokens":40,"cache_creation_tokens":5}}}`), "application/json")
require.ErrorIs(t, parseErr, fingerprint.ErrIncomplete)
assert.True(t, answer.HasUsage)
assert.EqualValues(t, 5, answer.CacheWriteTokens)
assert.Empty(t, answer.Text)
_, ok := detectionResponseUsage([]byte(`{"status":"completed"}`), "application/json")
assert.False(t, ok)
_, ok = detectionResponseUsage([]byte("data: {\"type\":\"response.output_text.delta\",\"delta\":\"hello\"}\n\n"), "text/event-stream")
assert.False(t, ok)
usage, ok := detectionResponseUsage([]byte("data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"usage\":{\"input_tokens\":100,\"output_tokens\":10}}}\n\n"), "text/event-stream")
require.True(t, ok)
assert.EqualValues(t, 100, usage.Get("input_tokens").Int())
expression := `p * 2 + cr * 0.2 + c * 10`
price := service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: expression, ExprHash: billingexpr.ExprHashString(expression), GroupRatio: 0.18, QuotaPerUnit: 500000, ExprVersion: 1}}
encoded, err := common.Marshal(price)
require.NoError(t, err)
quota, err := service.DetectionUsageQuota(string(encoded), 100, 10, 40)
require.NoError(t, err)
assert.Equal(t, 21, quota)
_, err = service.DetectionUsageQuota(string(encoded), 10, 20, 11)
assert.Error(t, err)
at := time.Date(2026, 9, 24, 9, 0, 0, 0, time.UTC)
price.Snapshot.ExprString = `hour("UTC") == 9 ? p * 2 : p * 100`
price.Snapshot.ExprHash = billingexpr.ExprHashString(price.Snapshot.ExprString)
price.Request.At = &at
encoded, err = common.Marshal(price)
require.NoError(t, err)
quota, err = service.DetectionUsageQuota(string(encoded), 100, 0, 0)
require.NoError(t, err)
assert.Equal(t, 18, quota)
}
func TestDetectionArtworkCannotNavigate(t *testing.T) {
cleaned := sanitizeDetectionHTML(`<meta http-equiv="refresh" content="0;url=https://evil.test"><script>alert(1)</script><a href="https://evil.test" onclick="alert(1)">x</a><iframe srcdoc="bad"></iframe><svg><circle cx="20" cy="20" r="10"/></svg><style>circle{fill:red}</style>`)
for _, bad := range []string{"evil.test", "onclick", "<iframe", "http-equiv"} {
assert.NotContains(t, cleaned, bad)
}
assert.Contains(t, cleaned, "<circle")
assert.Contains(t, cleaned, "<style>")
}
func TestDetectionPriceBoundIncludesIntermediateTiers(t *testing.T) {
bound, err := billingexpr.DetectionCostBound(`len < 100 ? p * 2 : len < 1000 ? p * 100 : p * 1`, 32768, 16384)
require.NoError(t, err)
assert.Equal(t, float64(3276800), bound)
_, err = billingexpr.DetectionCostBound(`p / (c + 1)`, 32768, 16384)
require.Error(t, err)
bound, err = billingexpr.DetectionCostBound(`len <= 272000 ? tier("standard", p * 10 + c * 50 + cr * 1 + cc * 12) : tier("long", p * 20 + c * 100 + cr * 2 + cc * 25)`, 32768, 16384)
require.NoError(t, err)
assert.Greater(t, bound, float64(16384*100))
}
func TestDetectionLatestPerChannelStatus(t *testing.T) {
now := time.Now().Unix()
for _, tc := range []struct {
name string
ids []int
evidence map[int]service.ChannelDetectionEvidence
want int
}{
{"successful fingerprint", []int{1}, map[int]service.ChannelDetectionEvidence{1: {At: now, Success: true}}, 1},
{"failed availability", []int{1}, map[int]service.ChannelDetectionEvidence{1: {At: now, Success: false}}, 0},
{"partial coverage", []int{1, 2}, map[int]service.ChannelDetectionEvidence{1: {At: now, Success: true}}, 2},
{"failure and unknown", []int{1, 2}, map[int]service.ChannelDetectionEvidence{1: {At: now, Success: false}}, -1},
{"expired success", []int{1}, map[int]service.ChannelDetectionEvidence{1: {At: now - 7201, Success: true}}, -1},
} {
t.Run(tc.name, func(t *testing.T) {
status, _ := service.DetectionFallbackStatus(tc.ids, tc.evidence, now)
assert.Equal(t, tc.want, status)
})
}
}
func TestDetectionMeowPaidRoundAndResume(t *testing.T) {
oldRedis := common.RedisEnabled
common.RedisEnabled = false
t.Cleanup(func() { common.RedisEnabled = oldRedis })
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
conn, err := db.DB()
require.NoError(t, err)
conn.SetMaxOpenConns(1)
old := model.DB
model.DB = db
t.Cleanup(func() { model.DB = old; _ = conn.Close() })
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.User{}))
data, err := os.ReadFile("../pkg/model_fingerprint/testdata/meow-reference.json")
require.NoError(t, err)
var fixtures []struct {
Samples []fingerprint.Sample `json:"samples"`
}
require.NoError(t, common.Unmarshal(data, &fixtures))
calls := 0
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
require.Less(t, calls, 32)
answer := fixtures[0].Samples[calls].Text
calls++
body, encodeErr := common.Marshal(map[string]any{"status": "completed", "usage": map[string]int{"input_tokens": 10, "output_tokens": 2}, "output": []any{map[string]any{"type": "message", "content": []any{map[string]string{"type": "output_text", "text": answer}}}}})
require.NoError(t, encodeErr)
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write(body)
}))
defer upstream.Close()
url := upstream.URL
channel := model.Channel{Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: fingerprint.Model}
require.NoError(t, db.Create(&channel).Error)
owner := model.User{Username: "meow-test", Quota: 1000, Status: common.UserStatusEnabled}
require.NoError(t, db.Create(&owner).Error)
price, err := common.Marshal(service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: "p+c", ExprHash: billingexpr.ExprHashString("p+c"), GroupRatio: 1, QuotaPerUnit: 500000, ExprVersion: 1}, Request: billingexpr.RequestInput{Body: []byte(`{"max_output_tokens":128,"reasoning":{"effort":"low"}}`)}})
require.NoError(t, err)
run := model.DetectionRun{ID: "meow-test", RequestKey: "meow-test", UserID: owner.Id, Kind: "benchmark", ChannelID: channel.Id, Status: "pending", Deadline: time.Now().Unix() + 60, Reserved: 320, Price: string(price)}
require.NoError(t, model.CreateDetection(&run))
processPaidFingerprint(context.Background(), run)
require.NoError(t, db.First(&run, "id = ?", run.ID).Error)
assert.Equal(t, "succeeded", run.Status)
assert.Equal(t, "normal", run.Quality)
assert.True(t, run.Settled)
assert.Equal(t, 32, calls)
assert.Equal(t, 192, run.Charged)
require.NoError(t, db.First(&owner, owner.Id).Error)
assert.Equal(t, 808, owner.Quota)
processPaidFingerprint(context.Background(), run)
assert.Equal(t, 32, calls, "resume must not resend or recharge")
require.NoError(t, db.First(&owner, owner.Id).Error)
assert.Equal(t, 808, owner.Quota)
}
func TestDetectionModelTracePaidRoundAndResume(t *testing.T) {
oldRedis := common.RedisEnabled
common.RedisEnabled = false
t.Cleanup(func() { common.RedisEnabled = oldRedis })
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
conn, err := db.DB()
require.NoError(t, err)
conn.SetMaxOpenConns(1)
old := model.DB
model.DB = db
t.Cleanup(func() { model.DB = old; _ = conn.Close() })
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.User{}))
data, err := os.ReadFile("../pkg/model_fingerprint/testdata/modeltrace-reference.json")
require.NoError(t, err)
var fixtures []struct {
Samples []fingerprint.Sample `json:"samples"`
}
require.NoError(t, common.Unmarshal(data, &fixtures))
calls := 0
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
require.Less(t, calls, 3)
answer := fixtures[1].Samples[calls].Text
calls++
body, encodeErr := common.Marshal(map[string]any{"status": "completed", "usage": map[string]int{"input_tokens": 10, "output_tokens": 2}, "output": []any{map[string]any{"type": "message", "content": []any{map[string]string{"type": "output_text", "text": answer}}}}})
require.NoError(t, encodeErr)
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write(body)
}))
defer upstream.Close()
url := upstream.URL
channel := model.Channel{Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: fingerprint.Model}
require.NoError(t, db.Create(&channel).Error)
owner := model.User{Username: "meow-test", Quota: 1000, Status: common.UserStatusEnabled}
require.NoError(t, db.Create(&owner).Error)
price, err := common.Marshal(service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: "p+c", ExprHash: billingexpr.ExprHashString("p+c"), GroupRatio: 1, QuotaPerUnit: 500000, ExprVersion: 1}, Request: billingexpr.RequestInput{Body: []byte(`{"max_output_tokens":4096,"reasoning":{"effort":"low"}}`)}})
require.NoError(t, err)
run := model.DetectionRun{ID: "meow-test", RequestKey: "meow-test", UserID: owner.Id, Kind: "fingerprint", ChannelID: channel.Id, Status: "pending", Deadline: time.Now().Unix() + 60, Reserved: 320, Price: string(price)}
require.NoError(t, model.CreateDetection(&run))
processPaidFingerprint(context.Background(), run)
require.NoError(t, db.First(&run, "id = ?", run.ID).Error)
assert.Equal(t, "succeeded", run.Status)
assert.Equal(t, "match", run.Quality)
assert.True(t, run.Settled)
assert.Equal(t, 3, calls)
assert.Equal(t, 18, run.Charged)
require.NoError(t, db.First(&owner, owner.Id).Error)
assert.Equal(t, 982, owner.Quota)
processPaidFingerprint(context.Background(), run)
assert.Equal(t, 3, calls, "resume must not resend or recharge")
require.NoError(t, db.First(&owner, owner.Id).Error)
assert.Equal(t, 982, owner.Quota)
}
func TestCandyAutomaticRoundAndResume(t *testing.T) {
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
t.Run(dialect, func(t *testing.T) {
oldRedis := common.RedisEnabled
common.RedisEnabled = false
t.Cleanup(func() { common.RedisEnabled = oldRedis })
var driver gorm.Dialector = sqlite.Open(":memory:")
if dialect == "mysql" {
if os.Getenv("TEST_MYSQL_DSN") == "" {
t.Skip("TEST_MYSQL_DSN unset")
}
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
}
if dialect == "postgres" {
if os.Getenv("TEST_POSTGRES_DSN") == "" {
t.Skip("TEST_POSTGRES_DSN unset")
}
driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
}
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("candy_%d_", time.Now().UnixNano())}})
require.NoError(t, err)
conn, err := db.DB()
require.NoError(t, err)
conn.SetMaxOpenConns(1)
oldType, oldLogType := common.MainDatabaseType(), common.LogDatabaseType()
common.SetDatabaseTypes(common.DatabaseType(dialect), common.DatabaseType(dialect))
t.Cleanup(func() { common.SetDatabaseTypes(oldType, oldLogType) })
old := model.DB
model.DB = db
t.Cleanup(func() {
for _, row := range []any{&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.User{}, &model.DetectionSchedule{}, &model.SystemTask{}} {
require.NoError(t, db.Migrator().DropTable(row))
}
model.DB = old
_ = conn.Close()
})
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.User{}, &model.DetectionSchedule{}, &model.SystemTask{}))
calls := 0
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
require.Less(t, calls, 1)
answer := "21"
raw, readErr := io.ReadAll(r.Body)
require.NoError(t, readErr)
assert.Contains(t, gjson.GetBytes(raw, "input.0.content").String(), "圆形 7 9 8")
assert.NotContains(t, gjson.GetBytes(raw, "input.0.content").String(), "21")
calls++
body, encodeErr := common.Marshal(map[string]any{"status": "completed", "usage": map[string]int{"input_tokens": 10, "output_tokens": 2}, "output": []any{map[string]any{"type": "message", "content": []any{map[string]string{"type": "output_text", "text": answer}}}}})
require.NoError(t, encodeErr)
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write(body)
}))
defer upstream.Close()
url := upstream.URL
channel := model.Channel{Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: fingerprint.Model, Group: "candy-group"}
require.NoError(t, db.Create(&channel).Error)
owner := model.User{Username: "meow-test", Quota: 1000, Status: common.UserStatusEnabled}
require.NoError(t, db.Create(&owner).Error)
price, err := common.Marshal(service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: "p+c", ExprHash: billingexpr.ExprHashString("p+c"), GroupRatio: 1, QuotaPerUnit: 500000, ExprVersion: 1}, Request: billingexpr.RequestInput{Body: []byte(`{"max_output_tokens":4096,"reasoning":{"effort":"low"}}`)}})
require.NoError(t, err)
run := model.DetectionRun{ID: "meow-test", RequestKey: "meow-test", Kind: "candy", CreatedAt: time.Now().Unix(), ChannelID: channel.Id, Status: "pending", Deadline: time.Now().Unix() + 60, Reserved: 320, Price: string(price)}
require.NoError(t, model.CreateDetection(&run))
require.NoError(t, model.EnsureDetectionSchedule(context.Background(), model.DetectionSchedule{ChannelID: channel.Id, Kind: "candy", IntervalMinutes: 30}))
processDetection(context.Background(), run)
require.NoError(t, db.First(&run, "id = ?", run.ID).Error)
assert.Equal(t, "succeeded", run.Status)
assert.Equal(t, "normal", run.Quality)
assert.True(t, run.Settled)
var clock model.DetectionSchedule
require.NoError(t, db.First(&clock, "channel_id = ? AND kind = ?", channel.Id, "candy").Error)
assert.Equal(t, "normal", clock.LastValid)
assert.Greater(t, clock.NextDueAt, run.FinishedAt)
assert.Equal(t, 1, calls)
assert.Equal(t, 0, run.Charged, "automatic tests do not charge a customer")
require.NoError(t, db.First(&owner, owner.Id).Error)
assert.Equal(t, 1000, owner.Quota)
processDetection(context.Background(), run)
assert.Equal(t, 1, calls, "resume must not resend or recharge")
require.NoError(t, db.First(&owner, owner.Id).Error)
assert.Equal(t, 1000, owner.Quota)
summaries, err := service.GroupDetectionSummaries(context.Background(), time.Now(), map[string]*service.AstraSamplingSummary{}, map[string]model.StatusPageGroup{})
require.NoError(t, err)
require.NotNil(t, summaries["candy-group"].Candy)
assert.Equal(t, 1, summaries["candy-group"].Candy.Normal)
assert.Equal(t, "21", summaries["candy-group"].Candy.History[0].Closest)
assert.Zero(t, summaries["candy-group"].Benchmark.Normal)
})
}
}
func TestCandyAnswers(t *testing.T) {
for _, tc := range []struct{ answer, quality string }{
{"21", "normal"}, {" 21\n", "normal"}, {"28", "degraded"}, {"121", "degraded"}, {"21或28", "degraded"}, {"答案不是21", "degraded"}, {"", "degraded"},
} {
t.Run(tc.answer, func(t *testing.T) { assert.Equal(t, tc.quality, service.CandyQuality(tc.answer)) })
}
}