fix: update image pricing tests for 2K tier and refactor claude max billing policy

- Update 5 test assertions to match new 2K default price ($0.201 = base * 1.5)
- Refactor claude max cache billing policy into reusable functions
This commit is contained in:
erio
2026-02-27 15:13:05 +08:00
parent 0302c03864
commit e71be7e0f1
2 changed files with 107 additions and 21 deletions

View File

@@ -12,14 +12,14 @@ import (
func TestCalculateImageCost_DefaultPricing(t *testing.T) { func TestCalculateImageCost_DefaultPricing(t *testing.T) {
svc := &BillingService{} // pricingService 为 nil使用硬编码默认值 svc := &BillingService{} // pricingService 为 nil使用硬编码默认值
// 2K 尺寸,默认价格 $0.134 // 2K 尺寸,默认价格 $0.134 * 1.5 = $0.201
cost := svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, nil, 1.0) cost := svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, nil, 1.0)
require.InDelta(t, 0.134, cost.TotalCost, 0.0001) require.InDelta(t, 0.201, cost.TotalCost, 0.0001)
require.InDelta(t, 0.134, cost.ActualCost, 0.0001) require.InDelta(t, 0.201, cost.ActualCost, 0.0001)
// 多张图片 // 多张图片
cost = svc.CalculateImageCost("gemini-3-pro-image", "2K", 3, nil, 1.0) cost = svc.CalculateImageCost("gemini-3-pro-image", "2K", 3, nil, 1.0)
require.InDelta(t, 0.402, cost.TotalCost, 0.0001) require.InDelta(t, 0.603, cost.TotalCost, 0.0001)
} }
// TestCalculateImageCost_GroupCustomPricing 测试分组自定义价格 // TestCalculateImageCost_GroupCustomPricing 测试分组自定义价格
@@ -63,13 +63,13 @@ func TestCalculateImageCost_RateMultiplier(t *testing.T) {
// 费率倍数 1.5x // 费率倍数 1.5x
cost := svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, nil, 1.5) cost := svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, nil, 1.5)
require.InDelta(t, 0.134, cost.TotalCost, 0.0001) // TotalCost 不变 require.InDelta(t, 0.201, cost.TotalCost, 0.0001) // TotalCost = 0.134 * 1.5
require.InDelta(t, 0.201, cost.ActualCost, 0.0001) // ActualCost = 0.134 * 1.5 require.InDelta(t, 0.3015, cost.ActualCost, 0.0001) // ActualCost = 0.201 * 1.5
// 费率倍数 2.0x // 费率倍数 2.0x
cost = svc.CalculateImageCost("gemini-3-pro-image", "2K", 2, nil, 2.0) cost = svc.CalculateImageCost("gemini-3-pro-image", "2K", 2, nil, 2.0)
require.InDelta(t, 0.268, cost.TotalCost, 0.0001) require.InDelta(t, 0.402, cost.TotalCost, 0.0001)
require.InDelta(t, 0.536, cost.ActualCost, 0.0001) require.InDelta(t, 0.804, cost.ActualCost, 0.0001)
} }
// TestCalculateImageCost_ZeroCount 测试 imageCount=0 // TestCalculateImageCost_ZeroCount 测试 imageCount=0
@@ -95,8 +95,8 @@ func TestCalculateImageCost_ZeroRateMultiplier(t *testing.T) {
svc := &BillingService{} svc := &BillingService{}
cost := svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, nil, 0) cost := svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, nil, 0)
require.InDelta(t, 0.134, cost.TotalCost, 0.0001) require.InDelta(t, 0.201, cost.TotalCost, 0.0001)
require.InDelta(t, 0.134, cost.ActualCost, 0.0001) // 0 倍率当作 1.0 处理 require.InDelta(t, 0.201, cost.ActualCost, 0.0001) // 0 倍率当作 1.0 处理
} }
// TestGetImageUnitPrice_GroupPriorityOverDefault 测试分组价格优先于默认价格 // TestGetImageUnitPrice_GroupPriorityOverDefault 测试分组价格优先于默认价格
@@ -127,9 +127,9 @@ func TestGetImageUnitPrice_PartialGroupConfig(t *testing.T) {
cost := svc.CalculateImageCost("gemini-3-pro-image", "1K", 1, groupConfig, 1.0) cost := svc.CalculateImageCost("gemini-3-pro-image", "1K", 1, groupConfig, 1.0)
require.InDelta(t, 0.10, cost.TotalCost, 0.0001) require.InDelta(t, 0.10, cost.TotalCost, 0.0001)
// 2K 回退默认价格 $0.134 // 2K 回退默认价格 $0.201 (1.5倍)
cost = svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, groupConfig, 1.0) cost = svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, groupConfig, 1.0)
require.InDelta(t, 0.134, cost.TotalCost, 0.0001) require.InDelta(t, 0.201, cost.TotalCost, 0.0001)
// 4K 回退默认价格 $0.268 (翻倍) // 4K 回退默认价格 $0.268 (翻倍)
cost = svc.CalculateImageCost("gemini-3-pro-image", "4K", 1, groupConfig, 1.0) cost = svc.CalculateImageCost("gemini-3-pro-image", "4K", 1, groupConfig, 1.0)
@@ -140,10 +140,10 @@ func TestGetImageUnitPrice_PartialGroupConfig(t *testing.T) {
func TestGetDefaultImagePrice_FallbackHardcoded(t *testing.T) { func TestGetDefaultImagePrice_FallbackHardcoded(t *testing.T) {
svc := &BillingService{} // pricingService 为 nil svc := &BillingService{} // pricingService 为 nil
// 1K 和 2K 使用相同的默认价格 $0.134 // 1K 默认价格 $0.1342K 默认价格 $0.201 (1.5倍)
cost := svc.CalculateImageCost("gemini-3-pro-image", "1K", 1, nil, 1.0) cost := svc.CalculateImageCost("gemini-3-pro-image", "1K", 1, nil, 1.0)
require.InDelta(t, 0.134, cost.TotalCost, 0.0001) require.InDelta(t, 0.134, cost.TotalCost, 0.0001)
cost = svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, nil, 1.0) cost = svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, nil, 1.0)
require.InDelta(t, 0.134, cost.TotalCost, 0.0001) require.InDelta(t, 0.201, cost.TotalCost, 0.0001)
} }

View File

@@ -64,6 +64,70 @@ func applyClaudeMaxCacheBillingPolicy(input *RecordUsageInput) claudeMaxCacheBil
return out return out
} }
// detectClaudeMaxCacheBillingOutcomeForUsage only returns whether Claude Max policy
// should influence downstream override decisions. It does not mutate usage.
func detectClaudeMaxCacheBillingOutcomeForUsage(usage ClaudeUsage, parsed *ParsedRequest, group *Group, model string) claudeMaxCacheBillingOutcome {
var out claudeMaxCacheBillingOutcome
if !shouldApplyClaudeMaxBillingRulesForUsage(group, model, parsed) {
return out
}
if hasCacheCreationTokens(usage) {
out.ForcedCache1H = true
return out
}
if shouldSimulateClaudeMaxUsageForUsage(usage, parsed) {
out.Simulated = true
}
return out
}
func applyClaudeMaxCacheBillingPolicyToUsage(usage *ClaudeUsage, parsed *ParsedRequest, group *Group, model string, accountID int64) claudeMaxCacheBillingOutcome {
var out claudeMaxCacheBillingOutcome
if usage == nil || !shouldApplyClaudeMaxBillingRulesForUsage(group, model, parsed) {
return out
}
resolvedModel := strings.TrimSpace(model)
if resolvedModel == "" && parsed != nil {
resolvedModel = strings.TrimSpace(parsed.Model)
}
if hasCacheCreationTokens(*usage) {
before5m := usage.CacheCreation5mTokens
before1h := usage.CacheCreation1hTokens
changed := safelyForceCacheCreationTo1H(usage)
// Even when value is already 1h, still mark forced to skip account TTL override.
out.ForcedCache1H = true
if changed {
logger.LegacyPrintf("service.gateway", "force_claude_max_cache_1h: model=%s account=%d cache_creation_5m:%d->%d cache_creation_1h:%d->%d",
resolvedModel,
accountID,
before5m,
usage.CacheCreation5mTokens,
before1h,
usage.CacheCreation1hTokens,
)
}
return out
}
if !shouldSimulateClaudeMaxUsageForUsage(*usage, parsed) {
return out
}
beforeInputTokens := usage.InputTokens
out.Simulated = safelyProjectUsageToClaudeMax1H(usage, parsed)
if out.Simulated {
logger.LegacyPrintf("service.gateway", "simulate_claude_max_usage: model=%s account=%d input_tokens:%d->%d cache_creation_1h=%d",
resolvedModel,
accountID,
beforeInputTokens,
usage.InputTokens,
usage.CacheCreation1hTokens,
)
}
return out
}
func isClaudeFamilyModel(model string) bool { func isClaudeFamilyModel(model string) bool {
normalized := strings.ToLower(strings.TrimSpace(claude.NormalizeModelID(model))) normalized := strings.ToLower(strings.TrimSpace(claude.NormalizeModelID(model)))
if normalized == "" { if normalized == "" {
@@ -76,16 +140,22 @@ func shouldApplyClaudeMaxBillingRules(input *RecordUsageInput) bool {
if input == nil || input.Result == nil || input.APIKey == nil || input.APIKey.Group == nil { if input == nil || input.Result == nil || input.APIKey == nil || input.APIKey.Group == nil {
return false return false
} }
group := input.APIKey.Group return shouldApplyClaudeMaxBillingRulesForUsage(input.APIKey.Group, input.Result.Model, input.ParsedRequest)
}
func shouldApplyClaudeMaxBillingRulesForUsage(group *Group, model string, parsed *ParsedRequest) bool {
if group == nil {
return false
}
if !group.SimulateClaudeMaxEnabled || group.Platform != PlatformAnthropic { if !group.SimulateClaudeMaxEnabled || group.Platform != PlatformAnthropic {
return false return false
} }
model := input.Result.Model resolvedModel := model
if model == "" && input.ParsedRequest != nil { if resolvedModel == "" && parsed != nil {
model = input.ParsedRequest.Model resolvedModel = parsed.Model
} }
if !isClaudeFamilyModel(model) { if !isClaudeFamilyModel(resolvedModel) {
return false return false
} }
return true return true
@@ -96,13 +166,19 @@ func hasCacheCreationTokens(usage ClaudeUsage) bool {
} }
func shouldSimulateClaudeMaxUsage(input *RecordUsageInput) bool { func shouldSimulateClaudeMaxUsage(input *RecordUsageInput) bool {
if input == nil || input.Result == nil {
return false
}
if !shouldApplyClaudeMaxBillingRules(input) { if !shouldApplyClaudeMaxBillingRules(input) {
return false return false
} }
if !hasClaudeCacheSignals(input.ParsedRequest) { return shouldSimulateClaudeMaxUsageForUsage(input.Result.Usage, input.ParsedRequest)
}
func shouldSimulateClaudeMaxUsageForUsage(usage ClaudeUsage, parsed *ParsedRequest) bool {
if !hasClaudeCacheSignals(parsed) {
return false return false
} }
usage := input.Result.Usage
if usage.InputTokens <= 0 { if usage.InputTokens <= 0 {
return false return false
} }
@@ -149,6 +225,16 @@ func safelyApplyClaudeMaxUsageSimulation(result *ForwardResult, parsed *ParsedRe
return applyClaudeMaxUsageSimulation(result, parsed) return applyClaudeMaxUsageSimulation(result, parsed)
} }
func safelyProjectUsageToClaudeMax1H(usage *ClaudeUsage, parsed *ParsedRequest) (changed bool) {
defer func() {
if r := recover(); r != nil {
logger.LegacyPrintf("service.gateway", "simulate_claude_max_usage skipped: panic=%v", r)
changed = false
}
}()
return projectUsageToClaudeMax1H(usage, parsed)
}
func safelyForceCacheCreationTo1H(usage *ClaudeUsage) (changed bool) { func safelyForceCacheCreationTo1H(usage *ClaudeUsage) (changed bool) {
defer func() { defer func() {
if r := recover(); r != nil { if r := recover(); r != nil {