|
- package controller
-
- import (
- "encoding/json"
- "net/http"
- "net/http/httptest"
- "testing"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/model"
- "github.com/glebarez/sqlite"
- "github.com/gin-gonic/gin"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- func setupControllerLogIdentityDB(t *testing.T) *gorm.DB {
- t.Helper()
-
- db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
- require.NoError(t, err)
-
- origLogDB := model.LOG_DB
- origSQLite := common.UsingSQLite
- origMySQL := common.UsingMySQL
- origPostgreSQL := common.UsingPostgreSQL
-
- model.LOG_DB = db
- common.UsingSQLite = true
- common.UsingMySQL = false
- common.UsingPostgreSQL = false
-
- require.NoError(t, db.AutoMigrate(&model.Log{}))
-
- t.Cleanup(func() {
- model.LOG_DB = origLogDB
- common.UsingSQLite = origSQLite
- common.UsingMySQL = origMySQL
- common.UsingPostgreSQL = origPostgreSQL
- })
-
- return db
- }
-
- func TestGetAllLogsPassesChatIDAndUpstreamIDFilters(t *testing.T) {
- db := setupControllerLogIdentityDB(t)
-
- require.NoError(t, db.Create(&model.Log{
- UserId: 1,
- Username: "alice",
- CreatedAt: 1714465001,
- Type: model.LogTypeConsume,
- Content: "usage-a",
- ModelName: "gpt-4o-mini",
- TokenName: "demo",
- RequestId: "req_a",
- ChatId: "chat_a",
- UpstreamId: "up_a",
- }).Error)
- require.NoError(t, db.Create(&model.Log{
- UserId: 1,
- Username: "alice",
- CreatedAt: 1714465002,
- Type: model.LogTypeConsume,
- Content: "usage-b",
- ModelName: "gpt-4o-mini",
- TokenName: "demo",
- RequestId: "req_b",
- ChatId: "chat_b",
- UpstreamId: "up_b",
- }).Error)
-
- gin.SetMode(gin.TestMode)
- router := gin.New()
- router.GET("/api/log", GetAllLogs)
-
- req := httptest.NewRequest(http.MethodGet, "/api/log?type=2&chat_id=chat_a&upstream_id=up_a", nil)
- w := httptest.NewRecorder()
- router.ServeHTTP(w, req)
-
- require.Equal(t, http.StatusOK, w.Code)
-
- var payload struct {
- Success bool `json:"success"`
- Data struct {
- Items []model.Log `json:"items"`
- } `json:"data"`
- }
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &payload))
- require.True(t, payload.Success)
- require.Len(t, payload.Data.Items, 1)
- require.Equal(t, "chat_a", payload.Data.Items[0].ChatId)
- require.Equal(t, "up_a", payload.Data.Items[0].UpstreamId)
- }
|