Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

111 строки
2.5 KiB

  1. package model
  2. import (
  3. "testing"
  4. "github.com/glebarez/sqlite"
  5. "github.com/stretchr/testify/assert"
  6. "github.com/stretchr/testify/require"
  7. "gorm.io/gorm"
  8. )
  9. func setupUserModelRateLimitDB(t *testing.T) {
  10. t.Helper()
  11. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  12. require.NoError(t, err)
  13. sqlDB, _ := db.DB()
  14. sqlDB.SetMaxOpenConns(1)
  15. origDB := DB
  16. DB = db
  17. require.NoError(t, db.AutoMigrate(&UserModelRateLimit{}))
  18. t.Cleanup(func() {
  19. DB = origDB
  20. sqlDB.Close()
  21. })
  22. }
  23. func TestGetUserModelRateLimits_Empty(t *testing.T) {
  24. setupUserModelRateLimitDB(t)
  25. list, err := GetUserModelRateLimits(1)
  26. require.NoError(t, err)
  27. assert.Empty(t, list)
  28. }
  29. func TestSetUserModelRateLimits_CreateAndOverwrite(t *testing.T) {
  30. setupUserModelRateLimitDB(t)
  31. // 首次写入 2 条
  32. items := []UserModelRateLimit{
  33. {UserId: 1, Model: "gpt-4", Rpm: 60},
  34. {UserId: 1, Model: "gpt-3.5-turbo", Rpm: 120},
  35. }
  36. err := SetUserModelRateLimits(1, items)
  37. require.NoError(t, err)
  38. list, err := GetUserModelRateLimits(1)
  39. require.NoError(t, err)
  40. assert.Len(t, list, 2)
  41. // 再覆盖写入 1 条(只保留 gpt-4,rpm 改为 30)
  42. items2 := []UserModelRateLimit{
  43. {UserId: 1, Model: "gpt-4", Rpm: 30},
  44. }
  45. err = SetUserModelRateLimits(1, items2)
  46. require.NoError(t, err)
  47. list2, err := GetUserModelRateLimits(1)
  48. require.NoError(t, err)
  49. assert.Len(t, list2, 1)
  50. assert.Equal(t, "gpt-4", list2[0].Model)
  51. assert.Equal(t, 30, list2[0].Rpm)
  52. }
  53. func TestSetUserModelRateLimits_EmptySlice(t *testing.T) {
  54. setupUserModelRateLimitDB(t)
  55. // 先写 1 条
  56. items := []UserModelRateLimit{
  57. {UserId: 2, Model: "claude-3-opus", Rpm: 10},
  58. }
  59. err := SetUserModelRateLimits(2, items)
  60. require.NoError(t, err)
  61. list, err := GetUserModelRateLimits(2)
  62. require.NoError(t, err)
  63. assert.Len(t, list, 1)
  64. // 再用空 slice 清除
  65. err = SetUserModelRateLimits(2, []UserModelRateLimit{})
  66. require.NoError(t, err)
  67. list2, err := GetUserModelRateLimits(2)
  68. require.NoError(t, err)
  69. assert.Empty(t, list2)
  70. }
  71. func TestGetUserModelRpm_Found(t *testing.T) {
  72. setupUserModelRateLimitDB(t)
  73. items := []UserModelRateLimit{
  74. {UserId: 3, Model: "gemini-pro", Rpm: 20},
  75. }
  76. err := SetUserModelRateLimits(3, items)
  77. require.NoError(t, err)
  78. rpm, found := GetUserModelRpm(3, "gemini-pro")
  79. assert.True(t, found)
  80. assert.Equal(t, 20, rpm)
  81. }
  82. func TestGetUserModelRpm_NotFound(t *testing.T) {
  83. setupUserModelRateLimitDB(t)
  84. rpm, found := GetUserModelRpm(999, "nonexistent-model")
  85. assert.False(t, found)
  86. assert.Equal(t, 0, rpm)
  87. }