Vous ne pouvez pas sélectionner plus de 25 sujets Les noms de sujets doivent commencer par une lettre ou un nombre, peuvent contenir des tirets ('-') et peuvent comporter jusqu'à 35 caractères.
 
 
 

198 lignes
5.1 KiB

  1. package model
  2. import (
  3. "net/http/httptest"
  4. "testing"
  5. "github.com/QuantumNous/new-api/common"
  6. relaycommon "github.com/QuantumNous/new-api/relay/common"
  7. "github.com/gin-gonic/gin"
  8. "github.com/glebarez/sqlite"
  9. "github.com/stretchr/testify/require"
  10. "gorm.io/gorm"
  11. )
  12. func setupLogIdentityDB(t *testing.T) *gorm.DB {
  13. t.Helper()
  14. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  15. require.NoError(t, err)
  16. origDB := DB
  17. origLogDB := LOG_DB
  18. origSQLite := common.UsingSQLite
  19. origMySQL := common.UsingMySQL
  20. origPostgres := common.UsingPostgreSQL
  21. DB = db
  22. LOG_DB = db
  23. common.UsingSQLite = true
  24. common.UsingMySQL = false
  25. common.UsingPostgreSQL = false
  26. initCol()
  27. require.NoError(t, db.AutoMigrate(&Log{}))
  28. t.Cleanup(func() {
  29. DB = origDB
  30. LOG_DB = origLogDB
  31. common.UsingSQLite = origSQLite
  32. common.UsingMySQL = origMySQL
  33. common.UsingPostgreSQL = origPostgres
  34. initCol()
  35. })
  36. return db
  37. }
  38. func TestLogPersistsChatIDAndUpstreamID(t *testing.T) {
  39. db := setupLogIdentityDB(t)
  40. entry := &Log{
  41. UserId: 1,
  42. Username: "alice",
  43. CreatedAt: 1714464000,
  44. Type: LogTypeConsume,
  45. Content: "usage",
  46. ModelName: "gpt-4o-mini",
  47. TokenName: "demo",
  48. ChatId: "chat_123",
  49. UpstreamId: "req_upstream_123",
  50. RequestId: "req_internal_123",
  51. }
  52. require.NoError(t, db.Create(entry).Error)
  53. var saved Log
  54. require.NoError(t, db.First(&saved).Error)
  55. require.Equal(t, "chat_123", saved.ChatId)
  56. require.Equal(t, "req_upstream_123", saved.UpstreamId)
  57. require.Equal(t, "req_internal_123", saved.RequestId)
  58. }
  59. func TestGetAllLogsFiltersByChatIDAndUpstreamID(t *testing.T) {
  60. db := setupLogIdentityDB(t)
  61. require.NoError(t, db.Create(&Log{
  62. UserId: 1,
  63. Username: "alice",
  64. CreatedAt: 1714464001,
  65. Type: LogTypeConsume,
  66. Content: "usage-a",
  67. ModelName: "gpt-4o-mini",
  68. TokenName: "demo",
  69. ChatId: "chat_a",
  70. UpstreamId: "up_a",
  71. RequestId: "req_a",
  72. }).Error)
  73. require.NoError(t, db.Create(&Log{
  74. UserId: 2,
  75. Username: "bob",
  76. CreatedAt: 1714464002,
  77. Type: LogTypeConsume,
  78. Content: "usage-b",
  79. ModelName: "gpt-4o-mini",
  80. TokenName: "demo",
  81. ChatId: "chat_b",
  82. UpstreamId: "up_b",
  83. RequestId: "req_b",
  84. }).Error)
  85. logs, total, err := GetAllLogs(LogTypeConsume, 0, 0, "", "", "", 0, 20, 0, "", "", "chat_a", "")
  86. require.NoError(t, err)
  87. require.Equal(t, int64(1), total)
  88. require.Len(t, logs, 1)
  89. require.Equal(t, "chat_a", logs[0].ChatId)
  90. logs, total, err = GetAllLogs(LogTypeConsume, 0, 0, "", "", "", 0, 20, 0, "", "", "", "up_b")
  91. require.NoError(t, err)
  92. require.Equal(t, int64(1), total)
  93. require.Len(t, logs, 1)
  94. require.Equal(t, "up_b", logs[0].UpstreamId)
  95. }
  96. func TestGetUserLogsFiltersByChatIDAndUpstreamID(t *testing.T) {
  97. db := setupLogIdentityDB(t)
  98. require.NoError(t, db.Create(&Log{
  99. UserId: 9,
  100. Username: "owner",
  101. CreatedAt: 1714464010,
  102. Type: LogTypeConsume,
  103. Content: "usage-owner-a",
  104. ModelName: "gpt-4o-mini",
  105. TokenName: "demo",
  106. ChatId: "chat-owner-a",
  107. UpstreamId: "up-owner-a",
  108. RequestId: "req-owner-a",
  109. }).Error)
  110. require.NoError(t, db.Create(&Log{
  111. UserId: 9,
  112. Username: "owner",
  113. CreatedAt: 1714464011,
  114. Type: LogTypeConsume,
  115. Content: "usage-owner-b",
  116. ModelName: "gpt-4o-mini",
  117. TokenName: "demo",
  118. ChatId: "chat-owner-b",
  119. UpstreamId: "up-owner-b",
  120. RequestId: "req-owner-b",
  121. }).Error)
  122. logs, total, err := GetUserLogs(9, LogTypeConsume, 0, 0, "", "", 0, 20, "", "", "chat-owner-b", "")
  123. require.NoError(t, err)
  124. require.Equal(t, int64(1), total)
  125. require.Len(t, logs, 1)
  126. require.Equal(t, "chat-owner-b", logs[0].ChatId)
  127. logs, total, err = GetUserLogs(9, LogTypeConsume, 0, 0, "", "", 0, 20, "", "", "", "up-owner-a")
  128. require.NoError(t, err)
  129. require.Equal(t, int64(1), total)
  130. require.Len(t, logs, 1)
  131. require.Equal(t, "up-owner-a", logs[0].UpstreamId)
  132. }
  133. func TestRecordConsumeLogPersistsContextChatIDAndUpstreamID(t *testing.T) {
  134. db := setupLogIdentityDB(t)
  135. gin.SetMode(gin.TestMode)
  136. w := httptest.NewRecorder()
  137. c, _ := gin.CreateTestContext(w)
  138. c.Set("username", "alice")
  139. c.Set(common.RequestIdKey, "req_internal_ctx")
  140. relaycommon.SetRelayChatID(c, "chat_ctx")
  141. relaycommon.SetRelayUpstreamID(c, "up_ctx")
  142. RecordConsumeLog(c, 1, RecordConsumeLogParams{
  143. ChannelId: 1,
  144. ModelName: "gpt-4o-mini",
  145. TokenName: "demo",
  146. Content: "usage",
  147. Quota: 10,
  148. })
  149. var saved Log
  150. require.NoError(t, db.First(&saved).Error)
  151. require.Equal(t, "chat_ctx", saved.ChatId)
  152. require.Equal(t, "up_ctx", saved.UpstreamId)
  153. }
  154. func TestRecordErrorLogPersistsContextChatIDAndUpstreamID(t *testing.T) {
  155. db := setupLogIdentityDB(t)
  156. gin.SetMode(gin.TestMode)
  157. w := httptest.NewRecorder()
  158. c, _ := gin.CreateTestContext(w)
  159. c.Set("username", "alice")
  160. c.Set(common.RequestIdKey, "req_internal_err")
  161. relaycommon.SetRelayChatID(c, "chat_err")
  162. relaycommon.SetRelayUpstreamID(c, "up_err")
  163. RecordErrorLog(c, 1, 2, "gpt-4o-mini", "demo", "upstream failed", 3, 4, false, "default", map[string]interface{}{})
  164. var saved Log
  165. require.NoError(t, db.First(&saved).Error)
  166. require.Equal(t, "chat_err", saved.ChatId)
  167. require.Equal(t, "up_err", saved.UpstreamId)
  168. }