Files
new-api/dto/billing_usage_test.go
T
Calcium-Ion c36418c863 feat: enhance text protocol conversion and advanced custom routing (#5825)
* refactor: consolidate relay protocol converters

* refactor relayconvert text converters

* feat: refine relay converters and advanced custom routing

* refactor: enhance logging and add thought signature handling for Gemini requests

* refactor: enhance channel cache and pricing endpoint handling for advanced custom models

* feat: preserve billing usage semantics

* feat: add protocol-aware billing usage

* Delete useless files

* chore: update action versions in workflow files

* chore: update Docker action versions in workflow files

* fix: harden billing usage settlement and hot-path route matching

- estimate Gemini completion tokens locally when billable usageMetadata is
  prompt-only but output content was received (e.g. client aborts the stream
  before the final chunk), and rebuild the attached billing_usage as estimated
  so settlement does not bill zero output tokens
- guard NewClaudeMessagesBillingUsage against all-zero ClaudeUsage, matching
  the OpenAI/Gemini constructors, so a zero billing_usage cannot override a
  non-zero top-level usage during settlement
- cache compiled advanced-custom route model regexes; they run on the request
  hot path and were recompiled per request
- move the effectiveBillingUsage remap to PostTextConsumeQuota only, and
  document that calculateTextQuotaSummary expects remapped usage
- document the updatePricingLock -> channelSyncLock lock ordering that
  InitChannelCache/CacheUpdateChannel rely on, and the aux-struct pitfall in
  GeminiChatResponse.UnmarshalJSON
2026-07-11 20:44:12 +08:00

90 lines
3.5 KiB
Go

package dto
import (
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestNewGeminiChatBillingUsageRequiresTokenContent(t *testing.T) {
require.Nil(t, NewGeminiChatBillingUsage(nil))
require.Nil(t, NewGeminiChatBillingUsage(&GeminiUsageMetadata{}))
billingUsage := NewGeminiChatBillingUsage(&GeminiUsageMetadata{PromptTokenCount: 1})
require.NotNil(t, billingUsage)
require.NotNil(t, billingUsage.GeminiUsageMetadata)
assert.Equal(t, BillingUsageSourceGeminiChat, billingUsage.Source)
assert.Equal(t, BillingUsageSemanticGemini, billingUsage.Semantic)
assert.False(t, billingUsage.Estimated)
}
func TestNewClaudeMessagesBillingUsageRequiresTokenContent(t *testing.T) {
require.Nil(t, NewClaudeMessagesBillingUsage(nil))
require.Nil(t, NewClaudeMessagesBillingUsage(&ClaudeUsage{}))
require.Nil(t, NewClaudeMessagesBillingUsage(&ClaudeUsage{CacheCreation: &ClaudeCacheCreationUsage{}}))
billingUsage := NewClaudeMessagesBillingUsage(&ClaudeUsage{InputTokens: 1})
require.NotNil(t, billingUsage)
require.NotNil(t, billingUsage.ClaudeUsage)
assert.Equal(t, BillingUsageSourceClaudeMessages, billingUsage.Source)
assert.Equal(t, BillingUsageSemanticAnthropic, billingUsage.Semantic)
cacheOnly := NewClaudeMessagesBillingUsage(&ClaudeUsage{
CacheCreation: &ClaudeCacheCreationUsage{Ephemeral5mInputTokens: 4},
})
require.NotNil(t, cacheOnly)
}
func TestNewOpenAIChatBillingUsageRequiresTokenContent(t *testing.T) {
require.Nil(t, NewOpenAIChatBillingUsage(nil))
require.Nil(t, NewOpenAIChatBillingUsage(&Usage{}))
billingUsage := NewOpenAIChatBillingUsage(&Usage{PromptTokens: 1})
require.NotNil(t, billingUsage)
require.NotNil(t, billingUsage.OpenAIUsage)
assert.Equal(t, BillingUsageSourceOAIChat, billingUsage.Source)
assert.Equal(t, BillingUsageSemanticOpenAI, billingUsage.Semantic)
assert.Equal(t, 1, billingUsage.OpenAIUsage.PromptTokens)
}
func TestNewEstimatedGeminiChatBillingUsage(t *testing.T) {
billingUsage := NewEstimatedGeminiChatBillingUsage(&Usage{
PromptTokens: 11,
CompletionTokens: 7,
})
require.NotNil(t, billingUsage)
require.NotNil(t, billingUsage.GeminiUsageMetadata)
assert.True(t, billingUsage.Estimated)
assert.Equal(t, 11, billingUsage.GeminiUsageMetadata.PromptTokenCount)
assert.Equal(t, 7, billingUsage.GeminiUsageMetadata.CandidatesTokenCount)
assert.Equal(t, 18, billingUsage.GeminiUsageMetadata.TotalTokenCount)
}
func TestBillingUsageJSONUsesProtocolNamedFields(t *testing.T) {
billingUsage := &BillingUsage{
OpenAIUsage: &Usage{PromptTokens: 1, BillingUsage: NewClaudeMessagesBillingUsage(&ClaudeUsage{InputTokens: 9})},
ClaudeUsage: &ClaudeUsage{InputTokens: 2, BillingUsage: NewOpenAIChatBillingUsage(&Usage{PromptTokens: 8})},
GeminiUsageMetadata: &GeminiUsageMetadata{PromptTokenCount: 3, BillingUsage: NewOpenAIChatBillingUsage(&Usage{PromptTokens: 7})},
}
data, err := common.Marshal(billingUsage)
require.NoError(t, err)
assert.Contains(t, string(data), `"openai_usage"`)
assert.Contains(t, string(data), `"claude_usage"`)
assert.Contains(t, string(data), `"gemini_usage_metadata"`)
assert.NotContains(t, string(data), `"usage":`)
assert.NotContains(t, string(data), `"usage_metadata"`)
clone := CloneBillingUsage(billingUsage)
require.NotNil(t, clone.OpenAIUsage)
require.NotNil(t, clone.ClaudeUsage)
require.NotNil(t, clone.GeminiUsageMetadata)
assert.Nil(t, clone.OpenAIUsage.BillingUsage)
assert.Nil(t, clone.ClaudeUsage.BillingUsage)
assert.Nil(t, clone.GeminiUsageMetadata.BillingUsage)
}