controller/log_export_test.go
Rama main · réplica de lectura del repositorio original.
package controller
import (
"context"
"fmt"
"net/http"
"net/http/httptest"
"os"
"strings"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/alicebob/miniredis/v2"
"github.com/gin-gonic/gin"
"github.com/go-redis/redis/v8"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
// callLogHandler runs a log handler as the given user and returns the JSON
// lines of the response.
func callLogHandler(t *testing.T, handler gin.HandlerFunc, userId int, role int, query string) []map[string]any {
t.Helper()
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodGet, "/api/log?"+query, nil)
ctx.Set("id", userId)
ctx.Set("role", role)
handler(ctx)
var lines []map[string]any
for line := range strings.SplitSeq(strings.TrimSpace(recorder.Body.String()), "\n") {
var value map[string]any
require.NoError(t, common.UnmarshalJsonStr(line, &value))
lines = append(lines, value)
}
return lines
}
func exportedRequestIds(lines []map[string]any) []any {
ids := make([]any, 0, len(lines))
for _, line := range lines[:len(lines)-1] {
ids = append(ids, line["request_id"])
}
return ids
}
// seedExportLogs stores logs of two users: three in the last hour and one
// from 400 days ago.
func seedExportLogs(t *testing.T, mainDB *gorm.DB, logDB *gorm.DB, now int64) {
t.Helper()
require.NoError(t, mainDB.Create(&model.Channel{Id: 9, Name: "primary"}).Error)
logs := []model.Log{
{UserId: 1, Username: "alice", CreatedAt: now - 3600, Type: model.LogTypeConsume, ModelName: "gpt-a", ChannelId: 9, Quota: 10, RequestId: "a-1", Other: `{"admin_info":{"use_channel":["9"]}}`},
{UserId: 2, Username: "bob", CreatedAt: now - 1800, Type: model.LogTypeConsume, ModelName: "gpt-b", ChannelId: 9, Quota: 20, RequestId: "b-1"},
{UserId: 1, Username: "alice", CreatedAt: now - 60, Type: model.LogTypeError, ModelName: "gpt-a", RequestId: "a-2"},
{UserId: 1, Username: "alice", CreatedAt: now - 400*24*3600, Type: model.LogTypeConsume, ModelName: "gpt-a", RequestId: "a-old"},
}
for i := range logs {
require.NoError(t, logDB.Create(&logs[i]).Error)
}
}
func useLogExportDatabases(t *testing.T, mainDB *gorm.DB, logDB *gorm.DB, mainType common.DatabaseType, logType common.DatabaseType) {
t.Helper()
previousDB, previousLogDB := model.DB, model.LOG_DB
previousMain, previousLog := common.MainDatabaseType(), common.LogDatabaseType()
previousRedis, previousMemory := common.RedisEnabled, common.MemoryCacheEnabled
model.DB, model.LOG_DB = mainDB, logDB
common.SetDatabaseTypes(mainType, logType)
common.RedisEnabled, common.MemoryCacheEnabled = false, false
t.Cleanup(func() {
model.DB, model.LOG_DB = previousDB, previousLogDB
common.SetDatabaseTypes(previousMain, previousLog)
common.RedisEnabled, common.MemoryCacheEnabled = previousRedis, previousMemory
})
}
func TestLogExportDatabaseMatrix(t *testing.T) {
gin.SetMode(gin.TestMode)
previousMaster := common.IsMasterNode
common.IsMasterNode = true
t.Cleanup(func() { common.IsMasterNode = previousMaster })
for _, tc := range []struct {
kind, env string
typ common.DatabaseType
}{
{"sqlite", "", common.DatabaseTypeSQLite},
{"mysql", "AUDIT_MYSQL_DSN", common.DatabaseTypeMySQL},
{"postgres", "AUDIT_POSTGRES_DSN", common.DatabaseTypePostgreSQL},
{"clickhouse", "AUDIT_CLICKHOUSE_DSN", common.DatabaseTypeClickHouse},
} {
t.Run(tc.kind, func(t *testing.T) {
dsn := os.Getenv(tc.env)
if tc.env != "" && dsn == "" {
t.Skip(tc.env + " is not configured")
}
db, isolatedDSN := newAuditTestDatabase(t, tc.kind, dsn)
mainDB, mainType := db, tc.typ
if tc.kind == "clickhouse" {
mainDB, _ = newAuditTestDatabase(t, "sqlite", "")
mainType = common.DatabaseTypeSQLite
}
useLogExportDatabases(t, mainDB, db, mainType, tc.typ)
require.NoError(t, mainDB.AutoMigrate(&model.Channel{}))
if tc.kind == "clickhouse" {
t.Setenv("LOG_SQL_DSN", isolatedDSN)
require.NoError(t, model.InitLogDB())
} else {
require.NoError(t, db.AutoMigrate(&model.Log{}))
}
versionSQL := "SELECT version()"
if tc.kind == "sqlite" {
versionSQL = "SELECT sqlite_version()"
}
var version string
require.NoError(t, model.LOG_DB.Raw(versionSQL).Scan(&version).Error)
t.Logf("log database: %s %s", tc.kind, version)
now := time.Now().Unix()
seedExportLogs(t, mainDB, model.LOG_DB, now)
query := fmt.Sprintf("start_timestamp=%d&end_timestamp=%d", now-2*24*3600, now+3600)
listed := callLogHandler(t, GetUserLogs, 1, common.RoleCommonUser, query)[0]
require.Equal(t, true, listed["success"])
assert.EqualValues(t, 2, listed["data"].(map[string]any)["total"])
listed = callLogHandler(t, GetAllLogs, 3, common.RoleAdminUser, query+"&username=alice")[0]
require.Equal(t, true, listed["success"])
assert.EqualValues(t, 2, listed["data"].(map[string]any)["total"])
own := callLogHandler(t, ExportUserLogs, 1, common.RoleCommonUser, query)
require.Len(t, own, 3)
assert.Equal(t, []any{"a-2", "a-1"}, exportedRequestIds(own))
assert.Equal(t, map[string]any{"success": true, "message": "", "data": map[string]any{"count": float64(2)}}, own[2])
assert.Equal(t, float64(1), own[0]["id"], "user exports carry list positions, not row ids")
assert.Empty(t, own[1]["channel_name"])
assert.NotContains(t, own[1]["other"], "admin_info")
all := callLogHandler(t, ExportAllLogs, 3, common.RoleAdminUser, query+"&model_name=gpt-%25")
require.Len(t, all, 4)
assert.Equal(t, []any{"a-2", "b-1", "a-1"}, exportedRequestIds(all))
assert.Equal(t, "primary", all[2]["channel_name"])
assert.Contains(t, all[2]["other"], "admin_info")
assert.Equal(t, true, all[3]["success"])
})
}
}
func TestLogExportRules(t *testing.T) {
gin.SetMode(gin.TestMode)
db, _ := newAuditTestDatabase(t, "sqlite", "")
useLogExportDatabases(t, db, db, common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Log{}, &model.QuotaData{}))
now := time.Now().Unix()
seedExportLogs(t, db, db, now)
logSetting := operation_setting.GetLogSetting()
previousSetting := *logSetting
t.Cleanup(func() { *logSetting = previousSetting })
recent := fmt.Sprintf("start_timestamp=%d&end_timestamp=%d", now-2*24*3600, now)
tooLong := fmt.Sprintf("start_timestamp=%d&end_timestamp=%d", now-366*24*3600, now)
t.Run("time range", func(t *testing.T) {
rejected := callLogHandler(t, ExportAllLogs, 3, common.RoleAdminUser, tooLong)
require.Len(t, rejected, 1)
assert.Equal(t, false, rejected[0]["success"])
assert.Equal(t, "The time range cannot exceed {{days}} days", rejected[0]["message_key"])
assert.Equal(t, map[string]any{"days": float64(365)}, rejected[0]["message_params"])
assert.Equal(t, false, callLogHandler(t, GetUserLogs, 1, common.RoleCommonUser, tooLong)[0]["success"])
assert.Equal(t, false, callLogHandler(t, GetLogsSelfStat, 1, common.RoleCommonUser, tooLong)[0]["success"])
assert.Equal(t, true, callLogHandler(t, GetUserLogs, 3, common.RoleAdminUser, tooLong)[0]["success"], "administrators search their own logs without a limit")
// The dashboard follows the same range instead of its old one-month limit.
fortyDays := fmt.Sprintf("start_timestamp=%d&end_timestamp=%d", now-40*24*3600, now)
for _, handler := range []gin.HandlerFunc{GetUserQuotaDates, GetUserFlowQuotaDates} {
assert.Equal(t, "The time range cannot exceed {{days}} days", callLogHandler(t, handler, 1, common.RoleCommonUser, tooLong)[0]["message_key"])
assert.Equal(t, true, callLogHandler(t, handler, 1, common.RoleCommonUser, fortyDays)[0]["success"])
}
// Without a start, a user's search reaches back the whole range.
listed := callLogHandler(t, GetUserLogs, 1, common.RoleCommonUser, "")[0]
assert.Equal(t, true, listed["success"])
assert.EqualValues(t, 2, listed["data"].(map[string]any)["total"])
logSetting.MaxRangeDays = 400
assert.Equal(t, true, callLogHandler(t, GetUserLogs, 1, common.RoleCommonUser, tooLong)[0]["success"])
logSetting.MaxRangeDays = previousSetting.MaxRangeDays
})
t.Run("user export switch", func(t *testing.T) {
logSetting.UserExportEnabled = false
t.Cleanup(func() { logSetting.UserExportEnabled = previousSetting.UserExportEnabled })
disabled := callLogHandler(t, ExportUserLogs, 1, common.RoleCommonUser, recent)
require.Len(t, disabled, 1)
assert.Equal(t, "Log export is disabled", disabled[0]["message_key"])
assert.Equal(t, true, callLogHandler(t, ExportUserLogs, 3, common.RoleAdminUser, recent)[0]["success"], "administrators export their own logs even when users cannot")
})
t.Run("one export per user", func(t *testing.T) {
release, ok, err := service.AcquireLogExportLock(1)
require.NoError(t, err)
require.True(t, ok)
busy := callLogHandler(t, ExportUserLogs, 1, common.RoleCommonUser, recent)
require.Len(t, busy, 1)
assert.Equal(t, "Another log export is running. Try again after it finishes.", busy[0]["message_key"])
other := callLogHandler(t, ExportUserLogs, 2, common.RoleCommonUser, recent)
assert.Equal(t, true, other[len(other)-1]["success"], "other users are not blocked")
release()
assert.Len(t, callLogHandler(t, ExportUserLogs, 1, common.RoleCommonUser, recent), 3)
})
}
func TestAcquireLogExportLock(t *testing.T) {
for _, useRedis := range []bool{false, true} {
t.Run(fmt.Sprintf("redis=%v", useRedis), func(t *testing.T) {
previousRedis, previousClient := common.RedisEnabled, common.RDB
common.RedisEnabled = useRedis
if useRedis {
server := miniredis.RunT(t)
common.RDB = redis.NewClient(&redis.Options{Addr: server.Addr()})
require.NoError(t, common.RDB.Ping(context.Background()).Err())
}
t.Cleanup(func() { common.RedisEnabled, common.RDB = previousRedis, previousClient })
release, ok, err := service.AcquireLogExportLock(1)
require.NoError(t, err)
require.True(t, ok)
_, ok, err = service.AcquireLogExportLock(1)
require.NoError(t, err)
assert.False(t, ok, "a second export of the same user must wait")
releaseOther, ok, err := service.AcquireLogExportLock(2)
require.NoError(t, err)
assert.True(t, ok)
releaseOther()
release()
releaseAgain, ok, err := service.AcquireLogExportLock(1)
require.NoError(t, err)
assert.True(t, ok, "a finished export frees the user")
releaseAgain()
})
}
}