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.
 
 
 

160 lines
3.9 KiB

  1. package controller
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "net/http"
  6. "net/http/httptest"
  7. "testing"
  8. "github.com/QuantumNous/new-api/common"
  9. "github.com/QuantumNous/new-api/model"
  10. "github.com/glebarez/sqlite"
  11. "github.com/gin-gonic/gin"
  12. "github.com/stretchr/testify/assert"
  13. "github.com/stretchr/testify/require"
  14. "gorm.io/gorm"
  15. )
  16. func setupRedemptionControllerDB(t *testing.T) *gorm.DB {
  17. t.Helper()
  18. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  19. require.NoError(t, err)
  20. sqlDB, _ := db.DB()
  21. sqlDB.SetMaxOpenConns(1)
  22. origDB := model.DB
  23. origLogDB := model.LOG_DB
  24. model.DB = db
  25. model.LOG_DB = db
  26. common.UsingSQLite = true
  27. common.RedisEnabled = false
  28. require.NoError(t, db.AutoMigrate(&model.User{}, &model.Redemption{}, &model.Log{}))
  29. t.Cleanup(func() {
  30. model.DB = origDB
  31. model.LOG_DB = origLogDB
  32. _ = sqlDB.Close()
  33. })
  34. return db
  35. }
  36. func setupRedemptionRouter() *gin.Engine {
  37. gin.SetMode(gin.TestMode)
  38. r := gin.New()
  39. r.Use(func(c *gin.Context) {
  40. c.Set("id", 1)
  41. c.Next()
  42. })
  43. g := r.Group("/api/redemption")
  44. g.POST("/", AddRedemption)
  45. g.PUT("/", UpdateRedemption)
  46. return r
  47. }
  48. func TestAddRedemptionStoresRemark(t *testing.T) {
  49. db := setupRedemptionControllerDB(t)
  50. router := setupRedemptionRouter()
  51. body, err := json.Marshal(map[string]interface{}{
  52. "name": "campaign-a",
  53. "remark": "admin only note",
  54. "quota": 500000,
  55. "count": 1,
  56. "expired_time": 0,
  57. })
  58. require.NoError(t, err)
  59. req := httptest.NewRequest(http.MethodPost, "/api/redemption/", bytes.NewReader(body))
  60. req.Header.Set("Content-Type", "application/json")
  61. w := httptest.NewRecorder()
  62. router.ServeHTTP(w, req)
  63. assert.Equal(t, http.StatusOK, w.Code)
  64. var rows []model.Redemption
  65. require.NoError(t, db.Find(&rows).Error)
  66. require.Len(t, rows, 1)
  67. assert.Equal(t, "campaign-a", rows[0].Name)
  68. assert.Equal(t, "admin only note", rows[0].Remark)
  69. }
  70. func TestUpdateRedemptionStoresRemark(t *testing.T) {
  71. db := setupRedemptionControllerDB(t)
  72. router := setupRedemptionRouter()
  73. row := model.Redemption{
  74. Id: 1,
  75. UserId: 1,
  76. Key: "update-remark-key",
  77. Name: "campaign-b",
  78. Remark: "before update",
  79. Status: common.RedemptionCodeStatusEnabled,
  80. Quota: 500000,
  81. CreatedTime: common.GetTimestamp(),
  82. }
  83. require.NoError(t, db.Create(&row).Error)
  84. body, err := json.Marshal(map[string]interface{}{
  85. "id": 1,
  86. "name": "campaign-b",
  87. "remark": "after update",
  88. "quota": 500000,
  89. "expired_time": 0,
  90. })
  91. require.NoError(t, err)
  92. req := httptest.NewRequest(http.MethodPut, "/api/redemption/", bytes.NewReader(body))
  93. req.Header.Set("Content-Type", "application/json")
  94. w := httptest.NewRecorder()
  95. router.ServeHTTP(w, req)
  96. assert.Equal(t, http.StatusOK, w.Code)
  97. var stored model.Redemption
  98. require.NoError(t, db.First(&stored, 1).Error)
  99. assert.Equal(t, "after update", stored.Remark)
  100. }
  101. func TestUpdateRedemptionStatusOnlyKeepsRemark(t *testing.T) {
  102. db := setupRedemptionControllerDB(t)
  103. router := setupRedemptionRouter()
  104. row := model.Redemption{
  105. Id: 1,
  106. UserId: 1,
  107. Key: "status-only-remark-key",
  108. Name: "campaign-c",
  109. Remark: "keep this remark",
  110. Status: common.RedemptionCodeStatusEnabled,
  111. Quota: 500000,
  112. CreatedTime: common.GetTimestamp(),
  113. }
  114. require.NoError(t, db.Create(&row).Error)
  115. body, err := json.Marshal(map[string]interface{}{
  116. "id": 1,
  117. "status": common.RedemptionCodeStatusDisabled,
  118. "remark": "should not overwrite",
  119. })
  120. require.NoError(t, err)
  121. req := httptest.NewRequest(http.MethodPut, "/api/redemption/?status_only=true", bytes.NewReader(body))
  122. req.Header.Set("Content-Type", "application/json")
  123. w := httptest.NewRecorder()
  124. router.ServeHTTP(w, req)
  125. assert.Equal(t, http.StatusOK, w.Code)
  126. var stored model.Redemption
  127. require.NoError(t, db.First(&stored, 1).Error)
  128. assert.Equal(t, common.RedemptionCodeStatusDisabled, stored.Status)
  129. assert.Equal(t, "keep this remark", stored.Remark)
  130. }