feat: publish v1.0.0-rc.36-custom.1

controller/relay_error_log_test.go

@@ -0,0 +1,99 @@
+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)
+	}
+}