97 lines
4.1 KiB
Go
97 lines
4.1 KiB
Go
package model
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"math"
|
|
"net/http/httptest"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/schema"
|
|
)
|
|
|
|
func TestPersonalGroupConfigurationDatabaseMatrix(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("DSN unset")
|
|
}
|
|
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
|
|
case "postgres":
|
|
if os.Getenv("TEST_POSTGRES_DSN") == "" {
|
|
t.Skip("DSN unset")
|
|
}
|
|
driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
|
|
}
|
|
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("personal_%d_", time.Now().UnixNano())}})
|
|
require.NoError(t, err)
|
|
old, redis := DB, common.RedisEnabled
|
|
DB, common.RedisEnabled = db, false
|
|
ratios, special := ratio_setting.GroupRatio2JSONString(), ratio_setting.GroupGroupRatio2JSONString()
|
|
rows := []any{&PersonalGroupRatio{}, &GroupMetadata{}}
|
|
t.Cleanup(func() {
|
|
for _, row := range rows {
|
|
require.NoError(t, db.Migrator().DropTable(row))
|
|
}
|
|
DB, common.RedisEnabled = old, redis
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(ratios))
|
|
require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString(special))
|
|
conn, _ := db.DB()
|
|
_ = conn.Close()
|
|
})
|
|
require.NoError(t, db.AutoMigrate(rows...))
|
|
require.NoError(t, db.AutoMigrate(rows...))
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString("{\"default\":1,\"vip\":2}"))
|
|
require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString("{\"default\":{\"vip\":0.7}}"))
|
|
makeCtx := func() *gin.Context {
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
|
|
require.NoError(t, LoadPersonalGroupRatios(c, 42))
|
|
return c
|
|
}
|
|
require.Equal(t, 0.7, EffectiveGroupRatio(makeCtx(), "default", "vip").GroupRatio)
|
|
require.Equal(t, 1.0, EffectiveGroupRatio(makeCtx(), "default", "default").GroupRatio)
|
|
require.NoError(t, SavePersonalGroupRatios(context.Background(), 42, map[string]float64{"vip": 0.2, "default": 0}))
|
|
inFlight := makeCtx()
|
|
require.Equal(t, 0.2, EffectiveGroupRatio(inFlight, "default", "vip").GroupRatio)
|
|
require.Equal(t, 0.0, EffectiveGroupRatio(inFlight, "default", "default").GroupRatio)
|
|
require.True(t, HasPersonalGroupRatio(inFlight, "vip"))
|
|
for _, bad := range []float64{-1, math.NaN(), math.Inf(1)} {
|
|
require.Error(t, SavePersonalGroupRatios(context.Background(), 42, map[string]float64{"vip": bad}))
|
|
}
|
|
require.NoError(t, SavePersonalGroupRatios(context.Background(), 42, map[string]float64{}))
|
|
require.Equal(t, 0.2, EffectiveGroupRatio(inFlight, "default", "vip").GroupRatio)
|
|
require.Equal(t, 0.7, EffectiveGroupRatio(makeCtx(), "default", "vip").GroupRatio)
|
|
require.False(t, HasPersonalGroupRatio(makeCtx(), "vip"))
|
|
require.NoError(t, SaveGroupMetadata(context.Background(), map[string]GroupMetadataValue{"vip": {Tags: []string{"官Key", "官Key", " 打野 ", "不稳定"}, HideRecentDegraded: true}}))
|
|
metadata, err := GetGroupMetadata(context.Background())
|
|
require.NoError(t, err)
|
|
require.Equal(t, []string{"官Key", "打野", "不稳定"}, metadata["vip"].Tags)
|
|
require.True(t, metadata["vip"].HideRecentDegraded)
|
|
require.NoError(t, db.AutoMigrate(rows...))
|
|
metadata, err = GetGroupMetadata(context.Background())
|
|
require.NoError(t, err)
|
|
require.Len(t, metadata["vip"].Tags, 3)
|
|
var u User
|
|
require.NoError(t, json.Unmarshal([]byte("{\"personal_group_ratios\":{\"vip\":0},\"ratios\":{\"vip\":0}}"), &u))
|
|
require.Equal(t, 0.7, EffectiveGroupRatio(makeCtx(), "default", "vip").GroupRatio)
|
|
})
|
|
}
|
|
}
|