new-api/model/log_identity_test.go

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 TestGetUserLogsIncludesTaskSettlementOutsidePage(t *testing.T) {
	db := setupLogIdentityDB(t)
	preconsume := &Log{
		UserId: 1, CreatedAt: 100, Type: LogTypeConsume, ModelName: "seedance",
		Other: `{"billing_mode":"matrix","billing_phase":"preconsume","task_id":"task_cross_page"}`,
	}
	settlement := &Log{
		UserId: 1, CreatedAt: 200, Type: LogTypeConsume, ModelName: "seedance",
		Other: `{"billing_mode":"matrix","billing_phase":"settlement","task_id":"task_cross_page"}`,
	}
	require.NoError(t, db.Create(preconsume).Error)
	require.NoError(t, db.Create(settlement).Error)

	logs, total, err := GetUserLogs(1, LogTypeConsume, 0, 0, "", "", 1, 1, "", "", "", "")
	require.NoError(t, err)
	require.Equal(t, int64(2), total)
	require.Len(t, logs, 2)
	require.Equal(t, "preconsume", logBillingPhase(t, logs[0]))
	require.Equal(t, "settlement", logBillingPhase(t, logs[1]))
}

func TestGetAllLogsIncludesTaskSettlementOutsidePage(t *testing.T) {
	db := setupLogIdentityDB(t)
	preconsume := &Log{
		UserId: 1, CreatedAt: 100, Type: LogTypeConsume, ModelName: "seedance",
		Other: `{"billing_mode":"matrix","billing_phase":"preconsume","task_id":"task_cross_page_admin"}`,
	}
	settlement := &Log{
		UserId: 1, CreatedAt: 200, Type: LogTypeConsume, ModelName: "seedance",
		Other: `{"billing_mode":"matrix","billing_phase":"settlement","task_id":"task_cross_page_admin"}`,
	}
	require.NoError(t, db.Create(preconsume).Error)
	require.NoError(t, db.Create(settlement).Error)

	logs, total, err := GetAllLogs(LogTypeConsume, 0, 0, "", "", "", 1, 1, 0, "", "", "", "")
	require.NoError(t, err)
	require.Equal(t, int64(2), total)
	require.Len(t, logs, 2)
	require.Equal(t, "preconsume", logBillingPhase(t, logs[0]))
	require.Equal(t, "settlement", logBillingPhase(t, logs[1]))
}

func logBillingPhase(t *testing.T, log *Log) string {
	t.Helper()
	other := map[string]any{}
	require.NoError(t, common.UnmarshalJsonStr(log.Other, &other))
	phase, _ := other["billing_phase"].(string)
	return phase
}

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