You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

122 lines
2.7 KiB

  1. package model
  2. import (
  3. "testing"
  4. "github.com/QuantumNous/new-api/common"
  5. "github.com/glebarez/sqlite"
  6. "github.com/stretchr/testify/assert"
  7. "github.com/stretchr/testify/require"
  8. "gorm.io/gorm"
  9. )
  10. func setupChannelPricingDB(t *testing.T) *gorm.DB {
  11. t.Helper()
  12. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  13. require.NoError(t, err)
  14. sqlDB, _ := db.DB()
  15. sqlDB.SetMaxOpenConns(1)
  16. origDB := DB
  17. DB = db
  18. common.UsingSQLite = true
  19. common.RedisEnabled = false
  20. require.NoError(t, db.AutoMigrate(&ChannelPricing{}))
  21. t.Cleanup(func() {
  22. DB = origDB
  23. sqlDB.Close()
  24. })
  25. return db
  26. }
  27. func TestCacheWriteThrough_Insert(t *testing.T) {
  28. setupChannelPricingDB(t)
  29. cp := &ChannelPricing{
  30. ModelName: "test-insert-model",
  31. ChannelId: 9001,
  32. QuotaType: QuotaTypeByTokens,
  33. ModelRatio: 1.0,
  34. CacheRatio: 0.8,
  35. ImageRatio: 1.2,
  36. }
  37. require.NoError(t, cp.Insert())
  38. defer cp.Delete()
  39. found, ok := GetEffectivePricing("test-insert-model", 9001)
  40. assert.True(t, ok)
  41. assert.Equal(t, 0.8, found.CacheRatio)
  42. assert.Equal(t, 1.2, found.ImageRatio)
  43. }
  44. func TestCacheWriteThrough_Update(t *testing.T) {
  45. setupChannelPricingDB(t)
  46. cp := &ChannelPricing{
  47. ModelName: "test-update-model",
  48. ChannelId: 9002,
  49. QuotaType: QuotaTypeByTokens,
  50. ModelRatio: 1.0,
  51. }
  52. require.NoError(t, cp.Insert())
  53. defer cp.Delete()
  54. cp.CacheRatio = 0.9
  55. cp.AudioRatio = 1.5
  56. require.NoError(t, cp.Update())
  57. found, ok := GetEffectivePricing("test-update-model", 9002)
  58. assert.True(t, ok)
  59. assert.Equal(t, 0.9, found.CacheRatio)
  60. assert.Equal(t, 1.5, found.AudioRatio)
  61. }
  62. func TestCacheWriteThrough_Delete(t *testing.T) {
  63. setupChannelPricingDB(t)
  64. cp := &ChannelPricing{
  65. ModelName: "test-delete-model",
  66. ChannelId: 9003,
  67. QuotaType: QuotaTypeByTokens,
  68. ModelRatio: 1.0,
  69. }
  70. require.NoError(t, cp.Insert())
  71. require.NoError(t, cp.Delete())
  72. found, ok := GetEffectivePricing("test-delete-model", 9003)
  73. assert.False(t, ok)
  74. assert.Nil(t, found)
  75. }
  76. func TestExtendedFields_DefaultZero(t *testing.T) {
  77. setupChannelPricingDB(t)
  78. cp := &ChannelPricing{
  79. ModelName: "test-default-model",
  80. ChannelId: 9004,
  81. QuotaType: QuotaTypeByTokens,
  82. ModelRatio: 1.0,
  83. }
  84. require.NoError(t, cp.Insert())
  85. defer cp.Delete()
  86. found, ok := GetEffectivePricing("test-default-model", 9004)
  87. assert.True(t, ok)
  88. assert.Equal(t, 0.0, found.CacheRatio)
  89. assert.Equal(t, 0.0, found.CacheCreationRatio)
  90. assert.Equal(t, 0.0, found.ImageRatio)
  91. assert.Equal(t, 0.0, found.AudioRatio)
  92. assert.Equal(t, 0.0, found.AudioCompletionRatio)
  93. }
  94. func TestGetEffectivePricing_NotFound(t *testing.T) {
  95. setupChannelPricingDB(t)
  96. found, ok := GetEffectivePricing("nonexistent-model-xyz", 99999)
  97. assert.False(t, ok)
  98. assert.Nil(t, found)
  99. }