100 lines
3.5 KiB
Go
100 lines
3.5 KiB
Go
package controller
|
|
|
|
import (
|
|
"errors"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/QuantumNous/new-api/relaykit/types"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
func TestProcessChannelErrorUsesSnapshotWithoutLeakingChannelMetadata(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
previousDB, previousLogDB := model.DB, model.LOG_DB
|
|
previousRedisEnabled := common.RedisEnabled
|
|
previousMainDatabaseType := common.MainDatabaseType()
|
|
previousLogDatabaseType := common.LogDatabaseType()
|
|
previousErrorLogEnabled := constant.ErrorLogEnabled
|
|
|
|
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
|
require.NoError(t, err)
|
|
sqlDB, err := database.DB()
|
|
require.NoError(t, err)
|
|
sqlDB.SetMaxOpenConns(1)
|
|
require.NoError(t, database.AutoMigrate(&model.User{}, &model.Log{}))
|
|
model.DB, model.LOG_DB = database, database
|
|
common.RedisEnabled = false
|
|
common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
|
|
constant.ErrorLogEnabled = true
|
|
t.Cleanup(func() {
|
|
model.DB, model.LOG_DB = previousDB, previousLogDB
|
|
common.RedisEnabled = previousRedisEnabled
|
|
common.SetDatabaseTypes(previousMainDatabaseType, previousLogDatabaseType)
|
|
constant.ErrorLogEnabled = previousErrorLogEnabled
|
|
require.NoError(t, sqlDB.Close())
|
|
})
|
|
|
|
require.NoError(t, database.Create(&model.User{Id: 7, Username: "log-owner", Group: "default"}).Error)
|
|
recorder := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(recorder)
|
|
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
|
ctx.Set("id", 7)
|
|
ctx.Set("username", "log-owner")
|
|
ctx.Set("token_name", "test-token")
|
|
ctx.Set("token_id", 11)
|
|
ctx.Set("original_model", "gpt-test")
|
|
ctx.Set("group", "default")
|
|
ctx.Set("channel_id", 202)
|
|
ctx.Set("channel_name", "mutable-context-channel")
|
|
ctx.Set("channel_type", 9)
|
|
ctx.Set("use_channel", []string{"101"})
|
|
common.SetContextKey(ctx, constant.ContextKeyRequestStartTime, time.Now().Add(-time.Second))
|
|
|
|
channelSnapshot := types.ChannelError{
|
|
ChannelId: 101,
|
|
ChannelType: 1,
|
|
ChannelName: "snapshot-channel",
|
|
AutoBan: false,
|
|
}
|
|
apiErr := types.NewOpenAIError(errors.New("upstream failed"), types.ErrorCodeBadResponseStatusCode, http.StatusBadGateway)
|
|
|
|
processChannelError(ctx, channelSnapshot, apiErr, nil)
|
|
|
|
var stored model.Log
|
|
require.NoError(t, database.First(&stored).Error)
|
|
assert.Equal(t, channelSnapshot.ChannelId, stored.ChannelId)
|
|
storedOther, err := common.StrToMap(stored.Other)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, float64(http.StatusBadGateway), storedOther["status_code"])
|
|
for _, key := range []string{"channel_id", "channel_name", "channel_type"} {
|
|
assert.NotContains(t, storedOther, key)
|
|
}
|
|
adminInfo, ok := storedOther["admin_info"].(map[string]any)
|
|
require.True(t, ok)
|
|
assert.Equal(t, []any{"101"}, adminInfo["use_channel"])
|
|
|
|
logs, total, err := model.GetUserLogs(7, model.LogTypeError, 0, 0, "", "", 0, 10, "", "", "")
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(1), total)
|
|
require.Len(t, logs, 1)
|
|
assert.Equal(t, channelSnapshot.ChannelId, logs[0].ChannelId)
|
|
assert.Empty(t, logs[0].ChannelName)
|
|
userOther, err := common.StrToMap(logs[0].Other)
|
|
require.NoError(t, err)
|
|
assert.NotContains(t, userOther, "admin_info")
|
|
for _, key := range []string{"channel_id", "channel_name", "channel_type"} {
|
|
assert.NotContains(t, userOther, key)
|
|
}
|
|
}
|