global_model_pricing_test.go 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  1. package service
  2. import (
  3. "context"
  4. "database/sql"
  5. "testing"
  6. "github.com/stretchr/testify/require"
  7. )
  8. type stubGlobalPricingRepo struct {
  9. rows []GlobalModelPrice
  10. }
  11. func (s *stubGlobalPricingRepo) ListEnabled(context.Context) ([]GlobalModelPrice, error) {
  12. return append([]GlobalModelPrice(nil), s.rows...), nil
  13. }
  14. func (s *stubGlobalPricingRepo) ListAll(context.Context) ([]GlobalModelPrice, error) {
  15. return append([]GlobalModelPrice(nil), s.rows...), nil
  16. }
  17. func (s *stubGlobalPricingRepo) Create(_ context.Context, in GlobalModelPriceInput) (*GlobalModelPrice, error) {
  18. out := GlobalModelPrice{ID: 1, ModelPattern: in.ModelPattern, BillingMode: in.BillingMode,
  19. InputPrice: in.InputPrice, OutputPrice: in.OutputPrice, Enabled: in.Enabled}
  20. s.rows = append(s.rows, out)
  21. return &out, nil
  22. }
  23. func (s *stubGlobalPricingRepo) GetByID(_ context.Context, id int64) (*GlobalModelPrice, error) {
  24. for i := range s.rows {
  25. if s.rows[i].ID == id {
  26. cp := s.rows[i]
  27. return &cp, nil
  28. }
  29. }
  30. return nil, sql.ErrNoRows
  31. }
  32. func (s *stubGlobalPricingRepo) Update(_ context.Context, id int64, in GlobalModelPriceInput) (*GlobalModelPrice, error) {
  33. return &GlobalModelPrice{ID: id, ModelPattern: in.ModelPattern}, nil
  34. }
  35. func (s *stubGlobalPricingRepo) Delete(context.Context, int64) error { return nil }
  36. func (s *stubGlobalPricingRepo) SetEnabled(_ context.Context, id int64, enabled bool) error {
  37. for i := range s.rows {
  38. if s.rows[i].ID == id {
  39. s.rows[i].Enabled = enabled
  40. return nil
  41. }
  42. }
  43. return sql.ErrNoRows
  44. }
  45. func fptr(v float64) *float64 { return &v }
  46. func TestValidateGlobalModelPriceInput(t *testing.T) {
  47. valid := GlobalModelPriceInput{ModelPattern: "gpt-5", InputPrice: fptr(1), OutputPrice: fptr(2)}
  48. require.NoError(t, ValidateGlobalModelPriceInput(valid))
  49. require.NoError(t, ValidateGlobalModelPriceInput(GlobalModelPriceInput{ModelPattern: "gpt-5*", OutputPrice: fptr(2)}))
  50. require.NoError(t, ValidateGlobalModelPriceInput(GlobalModelPriceInput{
  51. ModelPattern: "x", BillingMode: BillingModePerRequest, PerRequestPrice: fptr(0.01)}))
  52. bad := []GlobalModelPriceInput{
  53. {ModelPattern: ""},
  54. {ModelPattern: "*"},
  55. {ModelPattern: "*foo"},
  56. {ModelPattern: "a*b"},
  57. {ModelPattern: "has space"},
  58. {ModelPattern: "x", BillingMode: "hourly"},
  59. {ModelPattern: "x", InputPrice: fptr(-1)},
  60. {ModelPattern: "x"},
  61. {ModelPattern: "x", BillingMode: BillingModePerRequest},
  62. {ModelPattern: "x", BillingMode: BillingModeVideo, InputPrice: fptr(1)},
  63. }
  64. for i, in := range bad {
  65. require.Error(t, ValidateGlobalModelPriceInput(in), "case %d", i)
  66. }
  67. }
  68. func TestGlobalModelPricingMatch(t *testing.T) {
  69. svc := NewGlobalModelPricingService(&stubGlobalPricingRepo{rows: []GlobalModelPrice{
  70. {ID: 1, ModelPattern: "gpt-5", BillingMode: BillingModeToken, InputPrice: fptr(1), Enabled: true},
  71. {ID: 2, ModelPattern: "gpt-*", BillingMode: BillingModeToken, InputPrice: fptr(5), Enabled: true},
  72. {ID: 3, ModelPattern: "gpt-5-*", BillingMode: BillingModeToken, InputPrice: fptr(7), Enabled: true},
  73. }})
  74. // 精确优先于通配
  75. hit := svc.Match(context.Background(), "GPT-5")
  76. require.NotNil(t, hit)
  77. require.Equal(t, 1.0, *hit.InputPrice)
  78. // 最长前缀通配
  79. hit = svc.Match(context.Background(), "gpt-5-turbo")
  80. require.NotNil(t, hit)
  81. require.Equal(t, 7.0, *hit.InputPrice)
  82. hit = svc.Match(context.Background(), "gpt-4o")
  83. require.NotNil(t, hit)
  84. require.Equal(t, 5.0, *hit.InputPrice)
  85. // 未命中
  86. require.Nil(t, svc.Match(context.Background(), "claude-sonnet-4"))
  87. // nil 服务恒未命中
  88. var nilSvc *GlobalModelPricingService
  89. require.Nil(t, nilSvc.Match(context.Background(), "gpt-5"))
  90. }
  91. func TestModelPricingResolverGlobalFirst(t *testing.T) {
  92. repo := &stubGlobalPricingRepo{rows: []GlobalModelPrice{
  93. {ID: 1, ModelPattern: "gpt-5", BillingMode: BillingModeToken,
  94. InputPrice: fptr(10), OutputPrice: fptr(20), Enabled: true},
  95. }}
  96. global := NewGlobalModelPricingService(repo)
  97. r := NewModelPricingResolverWithGlobal(nil, &BillingService{}, global)
  98. resolved := r.Resolve(context.Background(), PricingInput{
  99. Model: "gpt-5",
  100. Group: &Group{ID: 9, ModelPricing: []ChannelModelPricing{{
  101. Models: []string{"gpt-5"}, InputPrice: fptr(1), OutputPrice: fptr(1),
  102. }}},
  103. })
  104. require.Equal(t, PricingSourceGlobal, resolved.Source)
  105. require.NotNil(t, resolved.BasePricing)
  106. require.Equal(t, 10.0, resolved.BasePricing.InputPricePerToken)
  107. require.Equal(t, 20.0, resolved.BasePricing.OutputPricePerToken)
  108. }
  109. func TestModelPricingResolverWithoutGlobalKeepsLegacy(t *testing.T) {
  110. r := NewModelPricingResolver(nil, &BillingService{})
  111. resolved := r.Resolve(context.Background(), PricingInput{Model: "gpt-5"})
  112. require.NotEqual(t, PricingSourceGlobal, resolved.Source)
  113. }