| package controller |
|
|
| import ( |
| "testing" |
|
|
| "github.com/QuantumNous/new-api/dto" |
| "github.com/QuantumNous/new-api/model" |
| "github.com/stretchr/testify/require" |
| ) |
|
|
| func TestNormalizeModelNames(t *testing.T) { |
| result := normalizeModelNames([]string{ |
| " gpt-4o ", |
| "", |
| "gpt-4o", |
| "gpt-4.1", |
| " ", |
| }) |
|
|
| require.Equal(t, []string{"gpt-4o", "gpt-4.1"}, result) |
| } |
|
|
| func TestMergeModelNames(t *testing.T) { |
| result := mergeModelNames( |
| []string{"gpt-4o", "gpt-4.1"}, |
| []string{"gpt-4.1", " gpt-4.1-mini ", "gpt-4o"}, |
| ) |
|
|
| require.Equal(t, []string{"gpt-4o", "gpt-4.1", "gpt-4.1-mini"}, result) |
| } |
|
|
| func TestSubtractModelNames(t *testing.T) { |
| result := subtractModelNames( |
| []string{"gpt-4o", "gpt-4.1", "gpt-4.1-mini"}, |
| []string{"gpt-4.1", "not-exists"}, |
| ) |
|
|
| require.Equal(t, []string{"gpt-4o", "gpt-4.1-mini"}, result) |
| } |
|
|
| func TestIntersectModelNames(t *testing.T) { |
| result := intersectModelNames( |
| []string{"gpt-4o", "gpt-4.1", "gpt-4.1", "not-exists"}, |
| []string{"gpt-4.1", "gpt-4o-mini", "gpt-4o"}, |
| ) |
|
|
| require.Equal(t, []string{"gpt-4o", "gpt-4.1"}, result) |
| } |
|
|
| func TestApplySelectedModelChanges(t *testing.T) { |
| t.Run("add and remove together", func(t *testing.T) { |
| result := applySelectedModelChanges( |
| []string{"gpt-4o", "gpt-4.1", "claude-3"}, |
| []string{"gpt-4.1-mini"}, |
| []string{"claude-3"}, |
| ) |
|
|
| require.Equal(t, []string{"gpt-4o", "gpt-4.1", "gpt-4.1-mini"}, result) |
| }) |
|
|
| t.Run("add wins when conflict with remove", func(t *testing.T) { |
| result := applySelectedModelChanges( |
| []string{"gpt-4o"}, |
| []string{"gpt-4.1"}, |
| []string{"gpt-4.1"}, |
| ) |
|
|
| require.Equal(t, []string{"gpt-4o", "gpt-4.1"}, result) |
| }) |
| } |
|
|
| func TestCollectPendingApplyUpstreamModelChanges(t *testing.T) { |
| settings := dto.ChannelOtherSettings{ |
| UpstreamModelUpdateLastDetectedModels: []string{" gpt-4o ", "gpt-4o", "gpt-4.1"}, |
| UpstreamModelUpdateLastRemovedModels: []string{" old-model ", "", "old-model"}, |
| } |
|
|
| pendingAddModels, pendingRemoveModels := collectPendingApplyUpstreamModelChanges(settings) |
|
|
| require.Equal(t, []string{"gpt-4o", "gpt-4.1"}, pendingAddModels) |
| require.Equal(t, []string{"old-model"}, pendingRemoveModels) |
| } |
|
|
| func TestNormalizeChannelModelMapping(t *testing.T) { |
| modelMapping := `{ |
| " alias-model ": " upstream-model ", |
| "": "invalid", |
| "invalid-target": "" |
| }` |
| channel := &model.Channel{ |
| ModelMapping: &modelMapping, |
| } |
|
|
| result := normalizeChannelModelMapping(channel) |
| require.Equal(t, map[string]string{ |
| "alias-model": "upstream-model", |
| }, result) |
| } |
|
|
| func TestCollectPendingUpstreamModelChangesFromModels_WithModelMapping(t *testing.T) { |
| pendingAddModels, pendingRemoveModels := collectPendingUpstreamModelChangesFromModels( |
| []string{"alias-model", "gpt-4o", "stale-model"}, |
| []string{"gpt-4o", "gpt-4.1", "mapped-target"}, |
| []string{"gpt-4.1"}, |
| map[string]string{ |
| "alias-model": "mapped-target", |
| }, |
| ) |
|
|
| require.Equal(t, []string{}, pendingAddModels) |
| require.Equal(t, []string{"stale-model"}, pendingRemoveModels) |
| } |
|
|
| func TestBuildUpstreamModelUpdateTaskNotificationContent_OmitOverflowDetails(t *testing.T) { |
| channelSummaries := make([]upstreamModelUpdateChannelSummary, 0, 12) |
| for i := 0; i < 12; i++ { |
| channelSummaries = append(channelSummaries, upstreamModelUpdateChannelSummary{ |
| ChannelName: "channel-" + string(rune('A'+i)), |
| AddCount: i + 1, |
| RemoveCount: i, |
| }) |
| } |
|
|
| content := buildUpstreamModelUpdateTaskNotificationContent( |
| 24, |
| 12, |
| 56, |
| 21, |
| 9, |
| []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}, |
| channelSummaries, |
| []string{ |
| "gpt-4.1", "gpt-4.1-mini", "o3", "o4-mini", "gemini-2.5-pro", "claude-3.7-sonnet", |
| "qwen-max", "deepseek-r1", "llama-3.3-70b", "mistral-large", "command-r-plus", "doubao-pro-32k", |
| "hunyuan-large", |
| }, |
| []string{ |
| "gpt-3.5-turbo", "claude-2.1", "gemini-1.5-pro", "mixtral-8x7b", "qwen-plus", "glm-4", |
| "yi-large", "moonshot-v1", "doubao-lite", |
| }, |
| ) |
|
|
| require.Contains(t, content, "其余 4 个渠道已省略") |
| require.Contains(t, content, "其余 1 个已省略") |
| require.Contains(t, content, "失败渠道 ID(展示 10/12)") |
| require.Contains(t, content, "其余 2 个已省略") |
| } |
|
|
| func TestShouldSendUpstreamModelUpdateNotification(t *testing.T) { |
| channelUpstreamModelUpdateNotifyState.Lock() |
| channelUpstreamModelUpdateNotifyState.lastNotifiedAt = 0 |
| channelUpstreamModelUpdateNotifyState.lastChangedChannels = 0 |
| channelUpstreamModelUpdateNotifyState.lastFailedChannels = 0 |
| channelUpstreamModelUpdateNotifyState.Unlock() |
|
|
| baseTime := int64(2000000) |
|
|
| require.True(t, shouldSendUpstreamModelUpdateNotification(baseTime, 6, 0)) |
| require.False(t, shouldSendUpstreamModelUpdateNotification(baseTime+3600, 6, 0)) |
| require.True(t, shouldSendUpstreamModelUpdateNotification(baseTime+3600, 7, 0)) |
| require.False(t, shouldSendUpstreamModelUpdateNotification(baseTime+7200, 7, 0)) |
| require.True(t, shouldSendUpstreamModelUpdateNotification(baseTime+8000, 0, 3)) |
| require.False(t, shouldSendUpstreamModelUpdateNotification(baseTime+9000, 0, 3)) |
| require.True(t, shouldSendUpstreamModelUpdateNotification(baseTime+10000, 0, 4)) |
| require.True(t, shouldSendUpstreamModelUpdateNotification(baseTime+90000, 7, 0)) |
| require.True(t, shouldSendUpstreamModelUpdateNotification(baseTime+90001, 0, 0)) |
| } |
|
|