mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-10-28 08:35:33 +08:00
sstar16
34003e9c2e
Replace wall-clock identifiers with an atomic sequence and test names so repeated Windows runs cannot share usage-cache entries. Verification: latest main baseline failed with cumulative counts; patched affinity tests passed -count=20; full service suite and go vet passed. Signed-off-by: sstar16 <3221386464@qq.com>
111 lines
3.4 KiB
Go
111 lines
3.4 KiB
Go
package service
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http/httptest"
|
|
"sync/atomic"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/relaykit/dto"
|
|
"github.com/QuantumNous/new-api/relaykit/types"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
var channelAffinityTestSeq atomic.Int64
|
|
|
|
func affinityTestKeys(t *testing.T) (string, string) {
|
|
t.Helper()
|
|
key := fmt.Sprintf("%s:%d", t.Name(), channelAffinityTestSeq.Add(1))
|
|
return key, key
|
|
}
|
|
|
|
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, keyFP := affinityTestKeys(t)
|
|
usingGroup := "default"
|
|
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, keyFP := affinityTestKeys(t)
|
|
usingGroup := "default"
|
|
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, keyFP := affinityTestKeys(t)
|
|
usingGroup := "default"
|
|
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)
|
|
}
|