openai_gateway_search_surcharge_test.go 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "testing"
  6. "time"
  7. "github.com/stretchr/testify/require"
  8. )
  9. func TestCalculateOpenAIRecordUsageCost_SearchIsAdditiveToTokens(t *testing.T) {
  10. t.Parallel()
  11. price := 10.0 // $10 / 1k searches → 100 searches = $1.0
  12. svc := &OpenAIGatewayService{
  13. billingService: newTestBillingService(),
  14. }
  15. apiKey := &APIKey{
  16. Group: &Group{
  17. SearchPricePer1k: &price,
  18. },
  19. }
  20. // claude-sonnet-4 fallback: Input $3/MTok, Output $15/MTok
  21. // 1000 in + 500 out → 0.003 + 0.0075 = 0.0105
  22. // + 100 searches → +1.0 → total 1.0105
  23. cost, err := svc.calculateOpenAIRecordUsageCost(
  24. context.Background(),
  25. &OpenAIForwardResult{SearchCount: 100},
  26. apiKey,
  27. []string{"claude-sonnet-4"},
  28. 1.0,
  29. 1.0,
  30. 1.0,
  31. 1.0,
  32. UsageTokens{InputTokens: 1000, OutputTokens: 500},
  33. "",
  34. boolPtr(false),
  35. time.Time{},
  36. )
  37. require.NoError(t, err)
  38. require.NotNil(t, cost)
  39. require.InDelta(t, 1.0105, cost.ActualCost, 1e-9)
  40. require.InDelta(t, 1.0105, cost.TotalCost, 1e-9)
  41. }
  42. func TestCalculateOpenAIRecordUsageCost_SearchOnlyWhenNoTokenPricing(t *testing.T) {
  43. t.Parallel()
  44. price := 10.0
  45. svc := &OpenAIGatewayService{
  46. billingService: newTestBillingService(),
  47. }
  48. apiKey := &APIKey{
  49. Group: &Group{SearchPricePer1k: &price},
  50. }
  51. // Empty model list: token path fails; search-only surcharge still bills.
  52. cost, err := svc.calculateOpenAIRecordUsageCost(
  53. context.Background(),
  54. &OpenAIForwardResult{SearchCount: 100},
  55. apiKey,
  56. nil,
  57. 1.0,
  58. 1.0,
  59. 1.0,
  60. 1.0,
  61. UsageTokens{},
  62. "",
  63. boolPtr(false),
  64. time.Time{},
  65. )
  66. require.NoError(t, err)
  67. require.NotNil(t, cost)
  68. require.InDelta(t, 1.0, cost.ActualCost, 1e-9)
  69. }
  70. func TestGroupMediaPricingLooksIncomplete_VideoModelPricesComplete(t *testing.T) {
  71. t.Parallel()
  72. require.True(t, groupMediaPricingLooksIncomplete(nil))
  73. require.True(t, groupMediaPricingLooksIncomplete(&Group{}))
  74. require.False(t, groupMediaPricingLooksIncomplete(&Group{
  75. VideoModelPrices: map[string]map[string]float64{
  76. "grok-imagine-video": {"720p": 0.1},
  77. },
  78. }))
  79. price := 10.0
  80. require.False(t, groupMediaPricingLooksIncomplete(&Group{SearchPricePer1k: &price}))
  81. require.False(t, groupMediaPricingLooksIncomplete(&Group{AudioRealtimePricePerMin: &price}))
  82. // Legacy video price alone still marks complete (existing path).
  83. require.False(t, groupMediaPricingLooksIncomplete(&Group{VideoPrice720P: &price}))
  84. }
  85. func TestCalculateOpenAIRecordUsageCost_TokenPricingErrorNotSwallowedBySearch(t *testing.T) {
  86. t.Parallel()
  87. price := 10.0
  88. svc := &OpenAIGatewayService{
  89. billingService: newTestBillingService(),
  90. }
  91. apiKey := &APIKey{
  92. Group: &Group{SearchPricePer1k: &price},
  93. }
  94. // Unknown model → token pricing fails; search must not replace that with $0/$search bill.
  95. cost, err := svc.calculateOpenAIRecordUsageCost(
  96. context.Background(),
  97. &OpenAIForwardResult{SearchCount: 100},
  98. apiKey,
  99. []string{"totally-unknown-model-xyz-no-pricing"},
  100. 1.0,
  101. 1.0,
  102. 1.0,
  103. 1.0,
  104. UsageTokens{InputTokens: 1000, OutputTokens: 500},
  105. "",
  106. boolPtr(false),
  107. time.Time{},
  108. )
  109. require.Error(t, err)
  110. require.Nil(t, cost)
  111. }