You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

95 rivejä
2.2 KiB

  1. package controller
  2. import (
  3. "encoding/json"
  4. "net/http"
  5. "net/http/httptest"
  6. "testing"
  7. "github.com/QuantumNous/new-api/common"
  8. "github.com/QuantumNous/new-api/model"
  9. "github.com/glebarez/sqlite"
  10. "github.com/gin-gonic/gin"
  11. "github.com/stretchr/testify/require"
  12. "gorm.io/gorm"
  13. )
  14. func setupControllerLogIdentityDB(t *testing.T) *gorm.DB {
  15. t.Helper()
  16. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  17. require.NoError(t, err)
  18. origLogDB := model.LOG_DB
  19. origSQLite := common.UsingSQLite
  20. origMySQL := common.UsingMySQL
  21. origPostgreSQL := common.UsingPostgreSQL
  22. model.LOG_DB = db
  23. common.UsingSQLite = true
  24. common.UsingMySQL = false
  25. common.UsingPostgreSQL = false
  26. require.NoError(t, db.AutoMigrate(&model.Log{}))
  27. t.Cleanup(func() {
  28. model.LOG_DB = origLogDB
  29. common.UsingSQLite = origSQLite
  30. common.UsingMySQL = origMySQL
  31. common.UsingPostgreSQL = origPostgreSQL
  32. })
  33. return db
  34. }
  35. func TestGetAllLogsPassesChatIDAndUpstreamIDFilters(t *testing.T) {
  36. db := setupControllerLogIdentityDB(t)
  37. require.NoError(t, db.Create(&model.Log{
  38. UserId: 1,
  39. Username: "alice",
  40. CreatedAt: 1714465001,
  41. Type: model.LogTypeConsume,
  42. Content: "usage-a",
  43. ModelName: "gpt-4o-mini",
  44. TokenName: "demo",
  45. RequestId: "req_a",
  46. ChatId: "chat_a",
  47. UpstreamId: "up_a",
  48. }).Error)
  49. require.NoError(t, db.Create(&model.Log{
  50. UserId: 1,
  51. Username: "alice",
  52. CreatedAt: 1714465002,
  53. Type: model.LogTypeConsume,
  54. Content: "usage-b",
  55. ModelName: "gpt-4o-mini",
  56. TokenName: "demo",
  57. RequestId: "req_b",
  58. ChatId: "chat_b",
  59. UpstreamId: "up_b",
  60. }).Error)
  61. gin.SetMode(gin.TestMode)
  62. router := gin.New()
  63. router.GET("/api/log", GetAllLogs)
  64. req := httptest.NewRequest(http.MethodGet, "/api/log?type=2&chat_id=chat_a&upstream_id=up_a", nil)
  65. w := httptest.NewRecorder()
  66. router.ServeHTTP(w, req)
  67. require.Equal(t, http.StatusOK, w.Code)
  68. var payload struct {
  69. Success bool `json:"success"`
  70. Data struct {
  71. Items []model.Log `json:"items"`
  72. } `json:"data"`
  73. }
  74. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &payload))
  75. require.True(t, payload.Success)
  76. require.Len(t, payload.Data.Items, 1)
  77. require.Equal(t, "chat_a", payload.Data.Items[0].ChatId)
  78. require.Equal(t, "up_a", payload.Data.Items[0].UpstreamId)
  79. }