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)) }