package model import ( "testing" "github.com/glebarez/sqlite" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gorm.io/gorm" ) func setupUserModelRateLimitDB(t *testing.T) { 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 require.NoError(t, db.AutoMigrate(&UserModelRateLimit{})) t.Cleanup(func() { DB = origDB sqlDB.Close() }) } func TestGetUserModelRateLimits_Empty(t *testing.T) { setupUserModelRateLimitDB(t) list, err := GetUserModelRateLimits(1) require.NoError(t, err) assert.Empty(t, list) } func TestSetUserModelRateLimits_CreateAndOverwrite(t *testing.T) { setupUserModelRateLimitDB(t) // 首次写入 2 条 items := []UserModelRateLimit{ {UserId: 1, Model: "gpt-4", Rpm: 60}, {UserId: 1, Model: "gpt-3.5-turbo", Rpm: 120}, } err := SetUserModelRateLimits(1, items) require.NoError(t, err) list, err := GetUserModelRateLimits(1) require.NoError(t, err) assert.Len(t, list, 2) // 再覆盖写入 1 条(只保留 gpt-4,rpm 改为 30) items2 := []UserModelRateLimit{ {UserId: 1, Model: "gpt-4", Rpm: 30}, } err = SetUserModelRateLimits(1, items2) require.NoError(t, err) list2, err := GetUserModelRateLimits(1) require.NoError(t, err) assert.Len(t, list2, 1) assert.Equal(t, "gpt-4", list2[0].Model) assert.Equal(t, 30, list2[0].Rpm) } func TestSetUserModelRateLimits_EmptySlice(t *testing.T) { setupUserModelRateLimitDB(t) // 先写 1 条 items := []UserModelRateLimit{ {UserId: 2, Model: "claude-3-opus", Rpm: 10}, } err := SetUserModelRateLimits(2, items) require.NoError(t, err) list, err := GetUserModelRateLimits(2) require.NoError(t, err) assert.Len(t, list, 1) // 再用空 slice 清除 err = SetUserModelRateLimits(2, []UserModelRateLimit{}) require.NoError(t, err) list2, err := GetUserModelRateLimits(2) require.NoError(t, err) assert.Empty(t, list2) } func TestGetUserModelRpm_Found(t *testing.T) { setupUserModelRateLimitDB(t) items := []UserModelRateLimit{ {UserId: 3, Model: "gemini-pro", Rpm: 20}, } err := SetUserModelRateLimits(3, items) require.NoError(t, err) rpm, found := GetUserModelRpm(3, "gemini-pro") assert.True(t, found) assert.Equal(t, 20, rpm) } func TestGetUserModelRpm_NotFound(t *testing.T) { setupUserModelRateLimitDB(t) rpm, found := GetUserModelRpm(999, "nonexistent-model") assert.False(t, found) assert.Equal(t, 0, rpm) }