|
- package model
-
- import (
- "testing"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/glebarez/sqlite"
- "github.com/stretchr/testify/assert"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- func setupUserChannelRatioDB(t *testing.T) *gorm.DB {
- t.Helper()
- db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
- require.NoError(t, err)
- sqlDB, _ := db.DB()
- sqlDB.SetMaxOpenConns(1)
-
- origDB := DB
- DB = db
- common.UsingSQLite = true
- common.RedisEnabled = false
-
- require.NoError(t, db.AutoMigrate(&UserChannelRatio{}))
-
- t.Cleanup(func() {
- DB = origDB
- sqlDB.Close()
- })
- return db
- }
-
- func TestUserChannelRatioCRUD(t *testing.T) {
- setupUserChannelRatioDB(t)
- LoadUserChannelRatioCache()
-
- ucr := &UserChannelRatio{
- UserId: 1,
- ModelName: "gpt-4o",
- ChannelId: 5,
- Ratio: 0.8,
- }
- err := ucr.Insert()
- require.NoError(t, err)
- assert.True(t, ucr.Id > 0)
-
- ratio := GetUserChannelRatio(1, "gpt-4o", 5)
- assert.Equal(t, 0.8, ratio)
-
- ratio = GetUserChannelRatio(999, "nonexistent", 999)
- assert.Equal(t, 1.0, ratio)
-
- ucr.Ratio = 1.2
- err = ucr.Update()
- require.NoError(t, err)
- ratio = GetUserChannelRatio(1, "gpt-4o", 5)
- assert.Equal(t, 1.2, ratio)
-
- list, err := GetUserChannelRatiosByUserId(1)
- require.NoError(t, err)
- assert.Len(t, list, 1)
- assert.Equal(t, "gpt-4o", list[0].ModelName)
-
- err = DeleteUserChannelRatioById(ucr.Id)
- require.NoError(t, err)
- ratio = GetUserChannelRatio(1, "gpt-4o", 5)
- assert.Equal(t, 1.0, ratio)
-
- list, err = GetUserChannelRatiosByUserId(1)
- require.NoError(t, err)
- assert.Len(t, list, 0)
- }
-
- func TestUserChannelRatioUniqueConstraint(t *testing.T) {
- setupUserChannelRatioDB(t)
- LoadUserChannelRatioCache()
-
- ucr1 := &UserChannelRatio{UserId: 1, ModelName: "gpt-4o", ChannelId: 5, Ratio: 0.8}
- err := ucr1.Insert()
- require.NoError(t, err)
-
- ucr2 := &UserChannelRatio{UserId: 1, ModelName: "gpt-4o", ChannelId: 5, Ratio: 1.0}
- err = ucr2.Insert()
- assert.Error(t, err)
-
- ucr3 := &UserChannelRatio{UserId: 1, ModelName: "gpt-4o", ChannelId: 6, Ratio: 1.5}
- err = ucr3.Insert()
- assert.NoError(t, err)
- }
-
- func TestLoadUserChannelRatioCache(t *testing.T) {
- setupUserChannelRatioDB(t)
-
- ucr1 := &UserChannelRatio{UserId: 1, ModelName: "gpt-4o", ChannelId: 5, Ratio: 0.8}
- ucr2 := &UserChannelRatio{UserId: 2, ModelName: "claude-3", ChannelId: 10, Ratio: 1.5}
- require.NoError(t, ucr1.Insert())
- require.NoError(t, ucr2.Insert())
-
- LoadUserChannelRatioCache()
-
- assert.Equal(t, 0.8, GetUserChannelRatio(1, "gpt-4o", 5))
- assert.Equal(t, 1.5, GetUserChannelRatio(2, "claude-3", 10))
- assert.Equal(t, 1.0, GetUserChannelRatio(3, "nonexistent", 99))
- }
|