new-api/service/channel_affinity_helpers_test.go

package service

import (
	"testing"

	"github.com/QuantumNous/new-api/dto"
	"github.com/QuantumNous/new-api/setting/operation_setting"
	"github.com/QuantumNous/new-api/types"
	"github.com/stretchr/testify/require"
)

func TestChannelAffinityMatchersIgnoreInvalidPatternsAndMatchCaseInsensitively(t *testing.T) {
	require.True(t, matchAnyRegexCached([]string{"[", `^gpt-[0-9]+$`}, "gpt-5"))
	require.False(t, matchAnyRegexCached([]string{"["}, "gpt-5"))
	require.True(t, matchAnyIncludeFold([]string{" Codex ", "other"}, "Mozilla CodexClient"))
	require.False(t, matchAnyIncludeFold([]string{""}, "CodexClient"))
}

func TestChannelAffinityKeyHelpersProduceStableSafeValues(t *testing.T) {
	rule := operation_setting.ChannelAffinityRule{Name: "rule", IncludeRuleName: true, IncludeUsingGroup: true}
	require.Equal(t, "rule:vip:key", buildChannelAffinityCacheKeySuffix(rule, "vip", "key"))
	require.Equal(t, "abcd...wxyz", buildChannelAffinityKeyHint("abcdefghijklmnopqrstuvwxwxyz"))
	require.Len(t, affinityFingerprint("tenant-123"), 8)
	require.Equal(t, "", channelAffinityUsageCacheEntryKey("", "vip", "fingerprint"))
	require.Equal(t, "rule\n\nfingerprint", channelAffinityUsageCacheEntryKey("rule", "", "fingerprint"))
}

func TestChannelAffinityUsageHelpersUseFallbackFields(t *testing.T) {
	usage := &dto.Usage{InputTokens: 3, OutputTokens: 5}
	require.Equal(t, 3, usagePromptTokens(usage))
	require.Equal(t, 5, usageCompletionTokens(usage))
	require.Equal(t, 8, usageTotalTokens(usage))
	require.Equal(t, cacheTokenRateModeCachedOverPrompt, cachedTokenRateModeByRelayFormat(types.RelayFormatOpenAI))
	require.Equal(t, cacheTokenRateModeCachedOverPromptPlusCached, cachedTokenRateModeByRelayFormat(types.RelayFormatClaude))
}