216 lines
13 KiB
Go
216 lines
13 KiB
Go
package controller
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/middleware"
|
|
"github.com/QuantumNous/new-api/model"
|
|
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
|
|
"github.com/QuantumNous/new-api/service"
|
|
"github.com/QuantumNous/new-api/setting"
|
|
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/schema"
|
|
"net/http/httptest"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestStatusPageDatabaseMatrix(t *testing.T) {
|
|
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
|
|
t.Run(dialect, func(t *testing.T) {
|
|
var driver gorm.Dialector
|
|
switch dialect {
|
|
case "sqlite":
|
|
driver = sqlite.Open(":memory:")
|
|
case "mysql":
|
|
if os.Getenv("TEST_MYSQL_DSN") == "" {
|
|
t.Skip("TEST_MYSQL_DSN unset")
|
|
}
|
|
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
|
|
case "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("status_%d_", time.Now().UnixNano())}})
|
|
require.NoError(t, err)
|
|
conn, err := db.DB()
|
|
require.NoError(t, err)
|
|
conn.SetMaxOpenConns(1)
|
|
oldDB, oldLog := model.DB, model.LOG_DB
|
|
oldType, oldLogType := common.MainDatabaseType(), common.LogDatabaseType()
|
|
oldRedis, oldSecret, oldSecure := common.RedisEnabled, common.SessionSecret, common.SessionCookieSecure
|
|
oldGroups := setting.UserUsableGroups2JSONString()
|
|
oldRatios := ratio_setting.GroupRatio2JSONString()
|
|
model.DB, model.LOG_DB = db, db
|
|
common.SetDatabaseTypes(common.DatabaseType(dialect), common.DatabaseType(dialect))
|
|
common.RedisEnabled, common.SessionSecret, common.SessionCookieSecure = false, "status-test-only", true
|
|
rows := []any{&model.GroupMetadata{}, &model.Option{}, &model.Channel{}, &model.Ability{}, &model.SystemTask{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.DetectionSchedule{}, &model.Log{}, &model.PerfMetric{}, &model.User{}, &model.UserSession{}}
|
|
require.NoError(t, db.AutoMigrate(rows...))
|
|
require.NoError(t, db.Table("audit_logs").AutoMigrate(&model.AuditLog{}))
|
|
t.Cleanup(func() {
|
|
require.NoError(t, db.Migrator().DropTable("audit_logs"))
|
|
for _, row := range rows {
|
|
require.NoError(t, db.Migrator().DropTable(row))
|
|
}
|
|
model.DB, model.LOG_DB = oldDB, oldLog
|
|
common.SetDatabaseTypes(oldType, oldLogType)
|
|
common.RedisEnabled, common.SessionSecret, common.SessionCookieSecure = oldRedis, oldSecret, oldSecure
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(oldGroups))
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(oldRatios))
|
|
groupCacheSnapshot.updated = time.Time{}
|
|
require.NoError(t, conn.Close())
|
|
})
|
|
var version string
|
|
q := "select version()"
|
|
if dialect == "sqlite" {
|
|
q = "select sqlite_version()"
|
|
}
|
|
require.NoError(t, db.Raw(q).Scan(&version).Error)
|
|
t.Log("database version:", version)
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"internal":"","shared":"","hidden":""}`))
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"internal":0.25,"shared":0,"hidden":1}`))
|
|
now := time.Now().Unix()
|
|
channels := []model.Channel{{Id: 1, Name: "private-channel", Key: "private-key", Group: "internal,shared", Models: fingerprint.Model, Status: 1}, {Id: 2, Group: "hidden", Models: fingerprint.Model, Status: 1}, {Id: 3, Group: "internal", Models: "disabled-model", Status: 2}}
|
|
require.NoError(t, db.Create(&channels).Error)
|
|
abilities := []model.Ability{{Group: "internal", Model: fingerprint.Model, ChannelId: 1, Enabled: true}, {Group: "shared", Model: fingerprint.Model, ChannelId: 1, Enabled: true}, {Group: "hidden", Model: fingerprint.Model, ChannelId: 2, Enabled: true}, {Group: "internal", Model: "disabled-model", ChannelId: 3, Enabled: true}, {Group: "internal", Model: "disabled-ability", ChannelId: 1}}
|
|
require.NoError(t, db.Create(&abilities).Error)
|
|
require.NoError(t, db.Create(&model.Log{Type: model.LogTypeConsume, TokenId: 1, ChannelId: 1, Group: "internal", ModelName: fingerprint.Model, Quota: 50, PromptTokens: 1000, CreatedAt: now - 10, Other: `{}`}).Error)
|
|
state, err := common.Marshal(service.AstraSamplingState{Rounds: []service.AstraSamplingRound{{ID: "old-round", ChannelID: 1, StartedAt: now - 20, FinishedAt: now - 10, Result: fingerprint.Result{Engine: fingerprint.Version, Status: "match", Closest: fingerprint.Model, ValidSamples: 3}}}})
|
|
require.NoError(t, err)
|
|
require.NoError(t, db.Create(&model.SystemTask{TaskID: "old-task", Type: service.AstraSamplingTaskType, Status: model.SystemTaskStatusSucceeded, CreatedAt: now - 20, State: string(state)}).Error)
|
|
runs := []model.DetectionRun{{ID: "old-work", RequestKey: "old-work", ChannelID: 1, GroupName: "internal", Kind: "benchmark", Status: "succeeded", Quality: "normal", CreatedAt: now - 20, Settled: true, HTML: "<h1>art</h1>", Charged: 17}, {ID: "in-flight", RequestKey: "in-flight", ChannelID: 1, Kind: "benchmark", Status: "running", CreatedAt: now - 10}, {ID: "hidden-art", RequestKey: "hidden-art", ChannelID: 2, Kind: "benchmark", Status: "succeeded", CreatedAt: now - 20, HTML: "hidden"}}
|
|
require.NoError(t, db.Create(&runs).Error)
|
|
tokens := map[int]string{}
|
|
for _, role := range []int{1, 10} {
|
|
require.NoError(t, db.Create(&model.User{Id: role, Username: fmt.Sprintf("role-%d", role), Role: role, Status: 1, AuthVersion: 1, AffCode: fmt.Sprintf("role-%d", role)}).Error)
|
|
sid := fmt.Sprintf("session-%d", role)
|
|
require.NoError(t, db.Create(&model.UserSession{SID: sid, UserID: role, Version: 1, UserAuthVersion: 1, Status: model.UserSessionStatusActive, RefreshHash: sid, ExpiresAt: now + 3600}).Error)
|
|
tokens[role], _, err = service.IssueAccessToken(service.AuthIdentity{UserID: role, SessionID: sid, UserAuthVersion: 1, SessionVersion: 1})
|
|
require.NoError(t, err)
|
|
}
|
|
engine := gin.New()
|
|
engine.GET("/api/uptime/status", GetGroupStatus)
|
|
engine.GET("/api/group-status", GetSiteGroupStatus)
|
|
engine.PUT("/api/uptime/group", middleware.AdminAuth(), middleware.SessionCookieOriginGuard(), UpdateStatusPageGroup)
|
|
engine.POST("/api/uptime/group/reset", middleware.AdminAuth(), middleware.SessionCookieOriginGuard(), UpdateStatusPageGroup)
|
|
engine.GET("/api/detections/artwork/:id", GetDetectionArtwork)
|
|
engine.GET("/api/group-status/artwork/:id", GetSiteDetectionArtwork)
|
|
request := func(method, path, body, token, origin string) *httptest.ResponseRecorder {
|
|
req := httptest.NewRequest(method, "https://status.test"+path, strings.NewReader(body))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
if token != "" {
|
|
req.Header.Set("Authorization", "Bearer "+token)
|
|
}
|
|
if origin != "" {
|
|
req.Header.Set("Origin", origin)
|
|
}
|
|
out := httptest.NewRecorder()
|
|
engine.ServeHTTP(out, req)
|
|
return out
|
|
}
|
|
for _, token := range []string{"", tokens[1]} {
|
|
assert.Contains(t, []int{401, 403}, request("POST", "/api/uptime/group/reset", `{"group":"internal"}`, token, "https://status.test").Code)
|
|
}
|
|
assert.Equal(t, 403, request("POST", "/api/uptime/group/reset", `{"group":"internal"}`, tokens[10], "https://evil.test").Code)
|
|
assert.Equal(t, 400, request("PUT", "/api/uptime/group", `{"group":"internal"}`, tokens[10], "https://status.test").Code)
|
|
assert.Equal(t, 404, request("POST", "/api/uptime/group/reset", `{"group":"missing"}`, tokens[10], "https://status.test").Code)
|
|
for _, body := range []string{`{"group":"internal","alias":"#688","hidden":false}`, `{"group":"hidden","alias":"secret-alias","hidden":true}`} {
|
|
res := request("PUT", "/api/uptime/group", body, tokens[10], "https://status.test")
|
|
require.Equal(t, 200, res.Code, res.Body.String())
|
|
}
|
|
var auditCount int64
|
|
require.NoError(t, db.Table("audit_logs").Count(&auditCount).Error)
|
|
assert.Greater(t, auditCount, int64(0))
|
|
groupCacheSnapshot.updated = time.Time{}
|
|
res := request("GET", "/api/uptime/status", "", "", "")
|
|
require.Equal(t, 200, res.Code, res.Body.String())
|
|
assert.NotContains(t, res.Body.String(), "model_costs")
|
|
for _, secret := range []string{"internal", "secret-alias", "private-channel", "private-key", "disabled-model", "disabled-ability"} {
|
|
assert.NotContains(t, res.Body.String(), secret)
|
|
}
|
|
assert.Contains(t, res.Body.String(), "#688")
|
|
assert.Contains(t, res.Body.String(), `"group_ratio":0.25`)
|
|
require.NoError(t, model.SaveGroupMetadata(context.Background(), map[string]model.GroupMetadataValue{"internal": {Tags: []string{"打野", "不稳定"}}}))
|
|
service.InvalidateGroupQualityCache()
|
|
visitorSite := request("GET", "/api/group-status", "", "", "")
|
|
require.Equal(t, 200, visitorSite.Code)
|
|
var siteData struct {
|
|
Data []UptimeGroupResult `json:"data"`
|
|
}
|
|
require.NoError(t, common.Unmarshal(visitorSite.Body.Bytes(), &siteData))
|
|
require.NotEmpty(t, siteData.Data)
|
|
for _, monitor := range siteData.Data[0].Monitors {
|
|
if monitor.Group == "internal" {
|
|
require.Equal(t, []string{"打野", "不稳定"}, monitor.Tags)
|
|
require.NotNil(t, monitor.ModelCosts[fingerprint.Model].USDPerMillion)
|
|
assert.Equal(t, int64(1), monitor.ModelCosts[fingerprint.Model].Requests)
|
|
}
|
|
if monitor.Group == "shared" {
|
|
assert.Empty(t, monitor.ModelCosts, "shared channel membership must not copy another group's billed usage")
|
|
}
|
|
}
|
|
assert.NotContains(t, request("GET", "/api/uptime/status", "", "", "").Body.String(), "model_costs")
|
|
assert.Contains(t, visitorSite.Body.String(), `"group":"hidden"`)
|
|
assert.Contains(t, visitorSite.Body.String(), `"name":"internal"`)
|
|
assert.NotContains(t, visitorSite.Body.String(), "#688")
|
|
site := request("GET", "/api/group-status", "", tokens[1], "")
|
|
require.Equal(t, 200, site.Code)
|
|
assert.Contains(t, site.Body.String(), `"group":"internal"`)
|
|
assert.Contains(t, site.Body.String(), `"name":"internal"`)
|
|
assert.NotContains(t, site.Body.String(), "#688")
|
|
assert.Contains(t, site.Body.String(), `"group":"hidden"`)
|
|
adminSite := request("GET", "/api/group-status", "", tokens[10], "")
|
|
assert.Contains(t, adminSite.Body.String(), `"group":"hidden"`)
|
|
assert.Equal(t, 404, request("GET", "/api/detections/artwork/hidden-art", "", "", "").Code)
|
|
assert.Equal(t, 200, request("GET", "/api/group-status/artwork/hidden-art", "", "", "").Code)
|
|
require.Equal(t, 200, request("POST", "/api/uptime/group/reset", `{"group":"internal"}`, tokens[10], "https://status.test").Code)
|
|
require.NoError(t, db.Model(&model.DetectionRun{}).Where("id = ?", "in-flight").Updates(map[string]any{"settled": true, "status": "succeeded", "quality": "normal", "finished_at": now + 10}).Error)
|
|
groups, err := currentGroupMonitors(context.Background(), true)
|
|
require.NoError(t, err)
|
|
byGroup := map[string]Monitor{}
|
|
for _, m := range groups[0].Monitors {
|
|
byGroup[m.Group] = m
|
|
}
|
|
assert.Empty(t, byGroup["internal"].Fingerprint.History)
|
|
assert.Nil(t, byGroup["internal"].Fingerprint.Rate)
|
|
assert.Empty(t, byGroup["internal"].Benchmark.History)
|
|
assert.Len(t, byGroup["shared"].Fingerprint.History, 1)
|
|
assert.Len(t, byGroup["shared"].Benchmark.History, 2)
|
|
assert.True(t, byGroup["hidden"].Hidden)
|
|
assert.Equal(t, []string{fingerprint.Model}, byGroup["internal"].Models)
|
|
require.NoError(t, model.SaveStatusPageGroup(context.Background(), "internal", &model.StatusPageGroup{Alias: "#689", Hidden: true}))
|
|
settings, err := model.GetStatusPageGroups(context.Background())
|
|
require.NoError(t, err)
|
|
assert.GreaterOrEqual(t, settings["internal"].ResetBefore, now)
|
|
assert.Equal(t, "#689", settings["internal"].Alias)
|
|
require.NoError(t, model.SaveStatusPageGroup(context.Background(), "shared", &model.StatusPageGroup{Hidden: true}))
|
|
assert.Equal(t, 404, request("GET", "/api/detections/artwork/old-work", "", "", "").Code)
|
|
var retained model.DetectionRun
|
|
require.NoError(t, db.First(&retained, "id = ?", "old-work").Error)
|
|
assert.Equal(t, 17, retained.Charged)
|
|
assert.Equal(t, "<h1>art</h1>", retained.HTML)
|
|
require.NoError(t, db.Create(&model.DetectionRun{ID: "new-work", RequestKey: "new-work", ChannelID: 1, Kind: "benchmark", Status: "succeeded", Quality: "normal", CreatedAt: settings["internal"].ResetBefore + 1, Settled: true}).Error)
|
|
groups, err = currentGroupMonitors(context.Background(), true)
|
|
require.NoError(t, err)
|
|
for _, m := range groups[0].Monitors {
|
|
if m.Group == "internal" {
|
|
require.Len(t, m.Benchmark.History, 1)
|
|
assert.Equal(t, "new-work", m.Benchmark.History[0].ID)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|