feat: sync deployed v1.0.0-rc.38-custom.42
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
This commit is contained in:
cneicy committed
1 parent
aa6c3f50f4
commit
ceddfd3d6b
828 files changed
+122253
-13702
No files matched your search
@@ -0,0 +1,225 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/mysql"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/schema"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestTokenRoutingResolution(t *testing.T) {
|
||||
p := TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []TokenGroupEntry{{"a", 0}, {"b", 0}, {"c", 2}}, ModelRules: []ModelRoutingRule{{Model: "model-x", Groups: []TokenGroupEntry{{"c", 0}, {"b", 0}}}}}
|
||||
allowed := func(string) bool { return true }
|
||||
require.NoError(t, p.Validate(6, allowed))
|
||||
require.Equal(t, []TokenGroupEntry{{"c", 0}, {"b", 0}, {"a", 1}}, p.Resolve("model-x", nil, allowed))
|
||||
require.Equal(t, []TokenGroupEntry{{"a", 0}, {"b", 0}, {"c", 1}}, p.Resolve("MODEL-X", nil, allowed))
|
||||
}
|
||||
|
||||
func TestTokenRoutingLegacy(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
token Token
|
||||
mode string
|
||||
groups []TokenGroupEntry
|
||||
}{
|
||||
{Token{Group: "auto"}, "inherit", []TokenGroupEntry{}},
|
||||
{Token{Group: ""}, "user", []TokenGroupEntry{}},
|
||||
{Token{Group: "a", AutoGroups: `["b","c"]`, CrossGroupRetry: true}, "custom", []TokenGroupEntry{{"a", 0}, {"b", 1}, {"c", 2}}},
|
||||
{Token{Group: "a", AutoGroups: `["b"]`, CrossGroupRetry: false}, "custom", []TokenGroupEntry{{"a", 0}}},
|
||||
{Token{Group: "a", AutoGroups: `[{"group":"b","priority":0},{"group":"a","priority":0}]`, CrossGroupRetry: true}, "custom", []TokenGroupEntry{{"b", 0}, {"a", 0}}},
|
||||
} {
|
||||
p, err := tc.token.LegacyRoutingPolicy()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tc.mode, p.Mode)
|
||||
require.Equal(t, tc.groups, p.Groups)
|
||||
require.Empty(t, p.ModelRules)
|
||||
}
|
||||
p := TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []TokenGroupEntry{{"a", 0}}, ModelRules: []ModelRoutingRule{{Model: "x", Groups: []TokenGroupEntry{{"b", 0}}}}}
|
||||
require.Error(t, p.Validate(6, func(string) bool { return true }))
|
||||
}
|
||||
|
||||
func TestRelayAttemptDatabaseMatrix(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":
|
||||
dsn := os.Getenv("TEST_MYSQL_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("TEST_MYSQL_DSN not configured")
|
||||
}
|
||||
driver = mysql.Open(dsn)
|
||||
case "postgres":
|
||||
dsn := os.Getenv("TEST_POSTGRES_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("TEST_POSTGRES_DSN not configured")
|
||||
}
|
||||
driver = postgres.Open(dsn)
|
||||
}
|
||||
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("attempt_%d_", time.Now().UnixNano())}})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
oldDB, oldLog, oldType, oldRedis := DB, LOG_DB, common.MainDatabaseType(), common.RedisEnabled
|
||||
DB = db
|
||||
LOG_DB = db
|
||||
common.SetMainDatabaseType(common.DatabaseType(dialect))
|
||||
common.RedisEnabled = false
|
||||
initCol()
|
||||
models := []any{&User{}, &Token{}, &Channel{}, &RelayAttempt{}, &Log{}, &QuotaData{}, &SubscriptionPlan{}, &UserSubscription{}}
|
||||
t.Cleanup(func() {
|
||||
for _, m := range models {
|
||||
require.NoError(t, db.Migrator().DropTable(m))
|
||||
}
|
||||
DB = oldDB
|
||||
LOG_DB = oldLog
|
||||
common.SetMainDatabaseType(oldType)
|
||||
common.RedisEnabled = oldRedis
|
||||
initCol()
|
||||
require.NoError(t, sqlDB.Close())
|
||||
})
|
||||
for range 2 {
|
||||
require.NoError(t, db.AutoMigrate(models...))
|
||||
}
|
||||
user := User{Username: "attempt-owner", Quota: 10000}
|
||||
require.NoError(t, db.Create(&user).Error)
|
||||
token := Token{UserId: user.Id, Key: "attempt-key", Status: common.TokenStatusEnabled, RemainQuota: 10000, Group: "a", CrossGroupRetry: true, AutoGroups: `["b"]`}
|
||||
require.NoError(t, db.Create(&token).Error)
|
||||
channel := Channel{Name: "attempt-channel", Status: common.ChannelStatusEnabled}
|
||||
require.NoError(t, db.Create(&channel).Error)
|
||||
// Upgrade a populated legacy token table, then run startup migrations twice.
|
||||
require.NoError(t, db.Migrator().DropColumn(&Token{}, "RoutingPolicy"))
|
||||
require.NoError(t, db.AutoMigrate(&Token{}))
|
||||
count, err := MigrateTokenRoutingPolicies()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, count)
|
||||
count, err = MigrateTokenRoutingPolicies()
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, count)
|
||||
var migrated Token
|
||||
require.NoError(t, db.First(&migrated, token.Id).Error)
|
||||
p, err := migrated.GetRoutingPolicy()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []TokenGroupEntry{{"a", 0}, {"b", 1}}, p.Groups)
|
||||
var total int
|
||||
for i, quota := range []int{120, 230, 0} {
|
||||
a := RelayAttempt{AttemptKey: fmt.Sprintf("request:%d", i), RequestID: "request", Number: i + 1, UserID: user.Id, TokenID: token.Id, ChannelID: channel.Id, Reserved: 300}
|
||||
require.NoError(t, ReserveRelayAttempt(&a, "wallet_only"))
|
||||
require.NoError(t, ReserveRelayAttempt(&a, "wallet_only"))
|
||||
require.NoError(t, CheckpointRelayAttempt(a.ID, quota, `{"total_tokens":10}`, "", "luna", "model_mismatch", "upstream", quota > 0, time.Now().Unix()))
|
||||
log := &Log{UserId: user.Id, Type: LogTypeConsume, RequestId: "request"}
|
||||
require.NoError(t, SettleRelayAttempt(a.ID, quota, log))
|
||||
require.NoError(t, SettleRelayAttempt(a.ID, quota+1, log))
|
||||
total += quota
|
||||
}
|
||||
require.NoError(t, DeliverRelayAttemptLogs())
|
||||
require.NoError(t, DeliverRelayAttemptLogs())
|
||||
require.NoError(t, db.First(&user, user.Id).Error)
|
||||
require.Equal(t, 10000-total, user.Quota)
|
||||
require.Equal(t, total, user.UsedQuota)
|
||||
require.Equal(t, 2, user.RequestCount)
|
||||
require.NoError(t, db.First(&token, token.Id).Error)
|
||||
require.Equal(t, 10000-total, token.RemainQuota)
|
||||
require.Equal(t, total, token.UsedQuota)
|
||||
require.NoError(t, db.First(&channel, channel.Id).Error)
|
||||
require.EqualValues(t, total, channel.UsedQuota)
|
||||
var logs int64
|
||||
require.NoError(t, db.Model(&Log{}).Count(&logs).Error)
|
||||
require.EqualValues(t, 3, logs)
|
||||
// A log-store outage never rolls back an already-committed financial fact.
|
||||
require.NoError(t, db.Model(&RelayAttempt{}).Where("request_id = ?", "request").Update("log_delivered", false).Error)
|
||||
require.NoError(t, DeliverRelayAttemptLogs())
|
||||
require.NoError(t, db.Model(&Log{}).Count(&logs).Error)
|
||||
require.EqualValues(t, 3, logs)
|
||||
plan := SubscriptionPlan{Title: "attempt-sub", QuotaResetPeriod: "never"}
|
||||
require.NoError(t, db.Create(&plan).Error)
|
||||
sub := UserSubscription{UserId: user.Id, PlanId: plan.Id, AmountTotal: 1000, Status: "active", EndTime: time.Now().Add(time.Hour).Unix()}
|
||||
require.NoError(t, db.Create(&sub).Error)
|
||||
subAttempt := RelayAttempt{AttemptKey: "sub-request", UserID: user.Id, TokenID: token.Id, ChannelID: channel.Id, Reserved: 300}
|
||||
require.NoError(t, ReserveRelayAttempt(&subAttempt, "subscription_only"))
|
||||
require.Equal(t, "subscription", subAttempt.Funding)
|
||||
require.NoError(t, SettleRelayAttempt(subAttempt.ID, 120, nil))
|
||||
require.NoError(t, SettleRelayAttempt(subAttempt.ID, 999, nil))
|
||||
require.NoError(t, db.First(&sub, sub.Id).Error)
|
||||
require.EqualValues(t, 120, sub.AmountUsed)
|
||||
require.NoError(t, db.First(&user, user.Id).Error)
|
||||
require.Equal(t, 10000-total, user.Quota)
|
||||
rejected := RelayAttempt{AttemptKey: "too-expensive", UserID: user.Id, TokenID: token.Id, Reserved: 20000}
|
||||
require.ErrorIs(t, ReserveRelayAttempt(&rejected, "wallet_only"), ErrAttemptQuota)
|
||||
require.NoError(t, db.First(&user, user.Id).Error)
|
||||
require.Equal(t, 10000-total, user.Quota)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRelayAttemptSQLiteConcurrentSettlementAndReserve(t *testing.T) {
|
||||
db, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "attempt.db")), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(8)
|
||||
oldDB, oldType, oldRedis := DB, common.MainDatabaseType(), common.RedisEnabled
|
||||
DB = db
|
||||
common.SetMainDatabaseType(common.DatabaseType("sqlite"))
|
||||
common.RedisEnabled = false
|
||||
initCol()
|
||||
t.Cleanup(func() {
|
||||
DB = oldDB
|
||||
common.SetMainDatabaseType(oldType)
|
||||
common.RedisEnabled = oldRedis
|
||||
initCol()
|
||||
_ = sqlDB.Close()
|
||||
})
|
||||
require.NoError(t, db.AutoMigrate(&User{}, &Token{}, &Channel{}, &RelayAttempt{}))
|
||||
user := User{Username: "concurrent", Quota: 10000}
|
||||
require.NoError(t, db.Create(&user).Error)
|
||||
token := Token{UserId: user.Id, Key: "concurrent-token", Status: common.TokenStatusEnabled, RemainQuota: 10000}
|
||||
require.NoError(t, db.Create(&token).Error)
|
||||
channel := Channel{Name: "concurrent"}
|
||||
require.NoError(t, db.Create(&channel).Error)
|
||||
var wg sync.WaitGroup
|
||||
failures := make(chan error, 16)
|
||||
for i := 0; i < 16; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
a := RelayAttempt{AttemptKey: fmt.Sprintf("concurrent:%d", i), UserID: user.Id, TokenID: token.Id, ChannelID: channel.Id, Reserved: 100}
|
||||
if err := ReserveRelayAttempt(&a, "wallet_only"); err != nil {
|
||||
failures <- err
|
||||
return
|
||||
}
|
||||
if err := SettleRelayAttempt(a.ID, 50, nil); err != nil {
|
||||
failures <- err
|
||||
return
|
||||
}
|
||||
if err := SettleRelayAttempt(a.ID, 50, nil); err != nil {
|
||||
failures <- err
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
close(failures)
|
||||
for err := range failures {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.NoError(t, db.First(&user, user.Id).Error)
|
||||
require.Equal(t, 9200, user.Quota)
|
||||
require.Equal(t, 800, user.UsedQuota)
|
||||
require.NoError(t, db.First(&token, token.Id).Error)
|
||||
require.Equal(t, 9200, token.RemainQuota)
|
||||
require.Equal(t, 800, token.UsedQuota)
|
||||
var count int64
|
||||
require.NoError(t, db.Model(&RelayAttempt{}).Where("state = ?", "settled").Count(&count).Error)
|
||||
require.EqualValues(t, 16, count)
|
||||
}
|
||||
Reference in new issue
Block a user