mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-03-29 23:28:36 +00:00
70 lines
2.0 KiB
Go
70 lines
2.0 KiB
Go
package service
|
|
|
|
import (
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func buildChannelAffinityTemplateContextForTest(meta channelAffinityMeta) *gin.Context {
|
|
rec := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(rec)
|
|
setChannelAffinityContext(ctx, meta)
|
|
return ctx
|
|
}
|
|
|
|
func TestApplyChannelAffinityOverrideTemplate_NoTemplate(t *testing.T) {
|
|
ctx := buildChannelAffinityTemplateContextForTest(channelAffinityMeta{
|
|
RuleName: "rule-no-template",
|
|
})
|
|
base := map[string]interface{}{
|
|
"temperature": 0.7,
|
|
}
|
|
|
|
merged, applied := ApplyChannelAffinityOverrideTemplate(ctx, base)
|
|
require.False(t, applied)
|
|
require.Equal(t, base, merged)
|
|
}
|
|
|
|
func TestApplyChannelAffinityOverrideTemplate_MergeTemplate(t *testing.T) {
|
|
ctx := buildChannelAffinityTemplateContextForTest(channelAffinityMeta{
|
|
RuleName: "rule-with-template",
|
|
ParamTemplate: map[string]interface{}{
|
|
"temperature": 0.2,
|
|
"top_p": 0.95,
|
|
},
|
|
UsingGroup: "default",
|
|
ModelName: "gpt-4.1",
|
|
RequestPath: "/v1/responses",
|
|
KeySourceType: "gjson",
|
|
KeySourcePath: "prompt_cache_key",
|
|
KeyHint: "abcd...wxyz",
|
|
KeyFingerprint: "abcd1234",
|
|
})
|
|
base := map[string]interface{}{
|
|
"temperature": 0.7,
|
|
"max_tokens": 2000,
|
|
}
|
|
|
|
merged, applied := ApplyChannelAffinityOverrideTemplate(ctx, base)
|
|
require.True(t, applied)
|
|
require.Equal(t, 0.2, merged["temperature"])
|
|
require.Equal(t, 0.95, merged["top_p"])
|
|
require.Equal(t, 2000, merged["max_tokens"])
|
|
require.Equal(t, 0.7, base["temperature"])
|
|
|
|
anyInfo, ok := ctx.Get(ginKeyChannelAffinityLogInfo)
|
|
require.True(t, ok)
|
|
info, ok := anyInfo.(map[string]interface{})
|
|
require.True(t, ok)
|
|
overrideInfoAny, ok := info["override_template"]
|
|
require.True(t, ok)
|
|
overrideInfo, ok := overrideInfoAny.(map[string]interface{})
|
|
require.True(t, ok)
|
|
require.Equal(t, true, overrideInfo["applied"])
|
|
require.Equal(t, "rule-with-template", overrideInfo["rule_name"])
|
|
require.EqualValues(t, 2, overrideInfo["param_override_keys"])
|
|
}
|