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.
 
 
 

80 lines
2.3 KiB

  1. package model
  2. import (
  3. "testing"
  4. "github.com/QuantumNous/new-api/common"
  5. "github.com/QuantumNous/new-api/constant"
  6. "github.com/QuantumNous/new-api/setting/ratio_setting"
  7. "github.com/glebarez/sqlite"
  8. "github.com/stretchr/testify/require"
  9. "gorm.io/gorm"
  10. )
  11. const testPricingModel = "test-pricing-global-model"
  12. func setupPricingTest(t *testing.T) {
  13. t.Helper()
  14. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  15. require.NoError(t, err)
  16. sqlDB, _ := db.DB()
  17. sqlDB.SetMaxOpenConns(1)
  18. origDB := DB
  19. origUsingSQLite := common.UsingSQLite
  20. DB = db
  21. common.UsingSQLite = true
  22. require.NoError(t, db.AutoMigrate(&Model{}, &Vendor{}, &Channel{}, &Ability{}))
  23. t.Cleanup(func() {
  24. DB = origDB
  25. common.UsingSQLite = origUsingSQLite
  26. sqlDB.Close()
  27. require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{}`))
  28. require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{}`))
  29. require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{}`))
  30. require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{}`))
  31. require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`))
  32. })
  33. }
  34. func TestGetPricingByModel_UsesGlobalDefaultOnly(t *testing.T) {
  35. setupPricingTest(t)
  36. require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingModel+`":10}`))
  37. require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingModel+`":3}`))
  38. require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingModel+`":0.5}`))
  39. require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingModel+`":0.75}`))
  40. require.NoError(t, (&Model{
  41. ModelName: testPricingModel,
  42. Status: 1,
  43. }).Insert())
  44. require.NoError(t, (&Vendor{
  45. Id: 1,
  46. Name: "test-vendor",
  47. SortOrder: 1,
  48. }).Insert())
  49. require.NoError(t, (&Channel{
  50. Id: 1,
  51. Type: constant.ChannelTypeOpenAI,
  52. Key: "k",
  53. Name: "channel-a",
  54. Status: common.ChannelStatusEnabled,
  55. Models: testPricingModel,
  56. Group: "default",
  57. }).Insert())
  58. RefreshPricing()
  59. pricing := GetPricingByModel(testPricingModel)
  60. require.NotNil(t, pricing)
  61. require.Equal(t, QuotaTypeByTokens, pricing.QuotaType)
  62. require.Equal(t, 10.0, pricing.ModelRatio)
  63. require.Equal(t, 3.0, pricing.CompletionRatio)
  64. require.Equal(t, 0.5, pricing.CacheRatio)
  65. require.Equal(t, 0.75, pricing.CacheCreationRatio)
  66. }