Spaces:
Runtime error
Runtime error
| 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)) | |
| } | |