|
- package model
-
- import (
- "net/http/httptest"
- "testing"
-
- "github.com/QuantumNous/new-api/common"
- relaycommon "github.com/QuantumNous/new-api/relay/common"
- "github.com/gin-gonic/gin"
- "github.com/glebarez/sqlite"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- func setupLogIdentityDB(t *testing.T) *gorm.DB {
- t.Helper()
-
- db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
- require.NoError(t, err)
-
- origDB := DB
- origLogDB := LOG_DB
- origSQLite := common.UsingSQLite
- origMySQL := common.UsingMySQL
- origPostgres := common.UsingPostgreSQL
-
- DB = db
- LOG_DB = db
- common.UsingSQLite = true
- common.UsingMySQL = false
- common.UsingPostgreSQL = false
- initCol()
-
- require.NoError(t, db.AutoMigrate(&Log{}))
-
- t.Cleanup(func() {
- DB = origDB
- LOG_DB = origLogDB
- common.UsingSQLite = origSQLite
- common.UsingMySQL = origMySQL
- common.UsingPostgreSQL = origPostgres
- initCol()
- })
-
- return db
- }
-
- func TestLogPersistsChatIDAndUpstreamID(t *testing.T) {
- db := setupLogIdentityDB(t)
-
- entry := &Log{
- UserId: 1,
- Username: "alice",
- CreatedAt: 1714464000,
- Type: LogTypeConsume,
- Content: "usage",
- ModelName: "gpt-4o-mini",
- TokenName: "demo",
- ChatId: "chat_123",
- UpstreamId: "req_upstream_123",
- RequestId: "req_internal_123",
- }
-
- require.NoError(t, db.Create(entry).Error)
-
- var saved Log
- require.NoError(t, db.First(&saved).Error)
- require.Equal(t, "chat_123", saved.ChatId)
- require.Equal(t, "req_upstream_123", saved.UpstreamId)
- require.Equal(t, "req_internal_123", saved.RequestId)
- }
-
- func TestGetAllLogsFiltersByChatIDAndUpstreamID(t *testing.T) {
- db := setupLogIdentityDB(t)
-
- require.NoError(t, db.Create(&Log{
- UserId: 1,
- Username: "alice",
- CreatedAt: 1714464001,
- Type: LogTypeConsume,
- Content: "usage-a",
- ModelName: "gpt-4o-mini",
- TokenName: "demo",
- ChatId: "chat_a",
- UpstreamId: "up_a",
- RequestId: "req_a",
- }).Error)
- require.NoError(t, db.Create(&Log{
- UserId: 2,
- Username: "bob",
- CreatedAt: 1714464002,
- Type: LogTypeConsume,
- Content: "usage-b",
- ModelName: "gpt-4o-mini",
- TokenName: "demo",
- ChatId: "chat_b",
- UpstreamId: "up_b",
- RequestId: "req_b",
- }).Error)
-
- logs, total, err := GetAllLogs(LogTypeConsume, 0, 0, "", "", "", 0, 20, 0, "", "", "chat_a", "")
- require.NoError(t, err)
- require.Equal(t, int64(1), total)
- require.Len(t, logs, 1)
- require.Equal(t, "chat_a", logs[0].ChatId)
-
- logs, total, err = GetAllLogs(LogTypeConsume, 0, 0, "", "", "", 0, 20, 0, "", "", "", "up_b")
- require.NoError(t, err)
- require.Equal(t, int64(1), total)
- require.Len(t, logs, 1)
- require.Equal(t, "up_b", logs[0].UpstreamId)
- }
-
- func TestGetUserLogsFiltersByChatIDAndUpstreamID(t *testing.T) {
- db := setupLogIdentityDB(t)
-
- require.NoError(t, db.Create(&Log{
- UserId: 9,
- Username: "owner",
- CreatedAt: 1714464010,
- Type: LogTypeConsume,
- Content: "usage-owner-a",
- ModelName: "gpt-4o-mini",
- TokenName: "demo",
- ChatId: "chat-owner-a",
- UpstreamId: "up-owner-a",
- RequestId: "req-owner-a",
- }).Error)
- require.NoError(t, db.Create(&Log{
- UserId: 9,
- Username: "owner",
- CreatedAt: 1714464011,
- Type: LogTypeConsume,
- Content: "usage-owner-b",
- ModelName: "gpt-4o-mini",
- TokenName: "demo",
- ChatId: "chat-owner-b",
- UpstreamId: "up-owner-b",
- RequestId: "req-owner-b",
- }).Error)
-
- logs, total, err := GetUserLogs(9, LogTypeConsume, 0, 0, "", "", 0, 20, "", "", "chat-owner-b", "")
- require.NoError(t, err)
- require.Equal(t, int64(1), total)
- require.Len(t, logs, 1)
- require.Equal(t, "chat-owner-b", logs[0].ChatId)
-
- logs, total, err = GetUserLogs(9, LogTypeConsume, 0, 0, "", "", 0, 20, "", "", "", "up-owner-a")
- require.NoError(t, err)
- require.Equal(t, int64(1), total)
- require.Len(t, logs, 1)
- require.Equal(t, "up-owner-a", logs[0].UpstreamId)
- }
-
- func TestRecordConsumeLogPersistsContextChatIDAndUpstreamID(t *testing.T) {
- db := setupLogIdentityDB(t)
-
- gin.SetMode(gin.TestMode)
- w := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(w)
- c.Set("username", "alice")
- c.Set(common.RequestIdKey, "req_internal_ctx")
- relaycommon.SetRelayChatID(c, "chat_ctx")
- relaycommon.SetRelayUpstreamID(c, "up_ctx")
-
- RecordConsumeLog(c, 1, RecordConsumeLogParams{
- ChannelId: 1,
- ModelName: "gpt-4o-mini",
- TokenName: "demo",
- Content: "usage",
- Quota: 10,
- })
-
- var saved Log
- require.NoError(t, db.First(&saved).Error)
- require.Equal(t, "chat_ctx", saved.ChatId)
- require.Equal(t, "up_ctx", saved.UpstreamId)
- }
-
- func TestRecordErrorLogPersistsContextChatIDAndUpstreamID(t *testing.T) {
- db := setupLogIdentityDB(t)
-
- gin.SetMode(gin.TestMode)
- w := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(w)
- c.Set("username", "alice")
- c.Set(common.RequestIdKey, "req_internal_err")
- relaycommon.SetRelayChatID(c, "chat_err")
- relaycommon.SetRelayUpstreamID(c, "up_err")
-
- RecordErrorLog(c, 1, 2, "gpt-4o-mini", "demo", "upstream failed", 3, 4, false, "default", map[string]interface{}{})
-
- var saved Log
- require.NoError(t, db.First(&saved).Error)
- require.Equal(t, "chat_err", saved.ChatId)
- require.Equal(t, "up_err", saved.UpstreamId)
- }
|