File size: 3,453 Bytes
daa8246 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | package service
import (
"fmt"
"net/http/httptest"
"testing"
"time"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP string) *gin.Context {
rec := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(rec)
setChannelAffinityContext(ctx, channelAffinityMeta{
CacheKey: fmt.Sprintf("test:%s:%s:%s", ruleName, usingGroup, keyFP),
TTLSeconds: 600,
RuleName: ruleName,
UsingGroup: usingGroup,
KeyFingerprint: keyFP,
})
return ctx
}
func TestObserveChannelAffinityUsageCacheByRelayFormat_ClaudeMode(t *testing.T) {
ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano())
usingGroup := "default"
keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano())
ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP)
usage := &dto.Usage{
PromptTokens: 100,
CompletionTokens: 40,
TotalTokens: 140,
PromptTokensDetails: dto.InputTokenDetails{
CachedTokens: 30,
},
}
ObserveChannelAffinityUsageCacheByRelayFormat(ctx, usage, types.RelayFormatClaude)
stats := GetChannelAffinityUsageCacheStats(ruleName, usingGroup, keyFP)
require.EqualValues(t, 1, stats.Total)
require.EqualValues(t, 1, stats.Hit)
require.EqualValues(t, 100, stats.PromptTokens)
require.EqualValues(t, 40, stats.CompletionTokens)
require.EqualValues(t, 140, stats.TotalTokens)
require.EqualValues(t, 30, stats.CachedTokens)
require.Equal(t, cacheTokenRateModeCachedOverPromptPlusCached, stats.CachedTokenRateMode)
}
func TestObserveChannelAffinityUsageCacheByRelayFormat_MixedMode(t *testing.T) {
ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano())
usingGroup := "default"
keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano())
ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP)
openAIUsage := &dto.Usage{
PromptTokens: 100,
PromptTokensDetails: dto.InputTokenDetails{
CachedTokens: 10,
},
}
claudeUsage := &dto.Usage{
PromptTokens: 80,
PromptTokensDetails: dto.InputTokenDetails{
CachedTokens: 20,
},
}
ObserveChannelAffinityUsageCacheByRelayFormat(ctx, openAIUsage, types.RelayFormatOpenAI)
ObserveChannelAffinityUsageCacheByRelayFormat(ctx, claudeUsage, types.RelayFormatClaude)
stats := GetChannelAffinityUsageCacheStats(ruleName, usingGroup, keyFP)
require.EqualValues(t, 2, stats.Total)
require.EqualValues(t, 2, stats.Hit)
require.EqualValues(t, 180, stats.PromptTokens)
require.EqualValues(t, 30, stats.CachedTokens)
require.Equal(t, cacheTokenRateModeMixed, stats.CachedTokenRateMode)
}
func TestObserveChannelAffinityUsageCacheByRelayFormat_UnsupportedModeKeepsEmpty(t *testing.T) {
ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano())
usingGroup := "default"
keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano())
ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP)
usage := &dto.Usage{
PromptTokens: 100,
PromptTokensDetails: dto.InputTokenDetails{
CachedTokens: 25,
},
}
ObserveChannelAffinityUsageCacheByRelayFormat(ctx, usage, types.RelayFormatGemini)
stats := GetChannelAffinityUsageCacheStats(ruleName, usingGroup, keyFP)
require.EqualValues(t, 1, stats.Total)
require.EqualValues(t, 1, stats.Hit)
require.EqualValues(t, 25, stats.CachedTokens)
require.Equal(t, "", stats.CachedTokenRateMode)
}
|