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