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) }