wans10andCaIon
58d4e9bd3b
* fix(billing): 异步任务退款时同步减少 used_quota 退款时仅恢复了 quota（剩余额度），但未同步减少 used_quota（已用额度）， 导致"总额度"（quota + used_quota）随退款次数持续虚增，超出用户实际充值金额。 修复三处退款路径： - RefundTaskQuota：任务失败完整退款 - RecalculateTaskQuota：差额结算退款分支 - controller/midjourney.go：Midjourney 任务失败退款 新增 model.UpdateUserUsedQuota 公开函数，仅调整 used_quota 不影响 request_count。 * fix(billing): 任务退款时同步扣减渠道 used_quota * fix(billing): complete async task refund accounting * style(model): group internal Midjourney fields --------- Co-authored-by: CaIon <i@caion.me>
401 lines
12 KiB
Go
401 lines
12 KiB
Go
package model
|
|
|
|
import (
|
|
"errors"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/relaykit/dto"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
func setupUserUpdateTestState(t *testing.T) {
|
|
t.Helper()
|
|
truncateTables(t)
|
|
require.NoError(t, DB.Exec("DELETE FROM users").Error)
|
|
|
|
oldRedisEnabled := common.RedisEnabled
|
|
oldBatchUpdateEnabled := common.BatchUpdateEnabled
|
|
common.RedisEnabled = false
|
|
common.BatchUpdateEnabled = false
|
|
t.Cleanup(func() {
|
|
common.RedisEnabled = oldRedisEnabled
|
|
common.BatchUpdateEnabled = oldBatchUpdateEnabled
|
|
})
|
|
}
|
|
|
|
func createUserBindTestUser(t *testing.T) User {
|
|
t.Helper()
|
|
user := User{
|
|
Username: "bind-test-user",
|
|
Password: "unused-password-hash",
|
|
Role: common.RoleCommonUser,
|
|
Status: common.UserStatusEnabled,
|
|
Group: "default",
|
|
AuthVersion: 1,
|
|
AffCode: "bind-test-aff-code",
|
|
}
|
|
require.NoError(t, DB.Create(&user).Error)
|
|
return user
|
|
}
|
|
|
|
func TestUserUpdateDoesNotOverwriteConcurrentAccountingOrTokenChanges(t *testing.T) {
|
|
setupUserUpdateTestState(t)
|
|
|
|
user := User{
|
|
Id: 1,
|
|
Username: "quota-race-user",
|
|
Password: "password",
|
|
DisplayName: "before",
|
|
Status: common.UserStatusEnabled,
|
|
Quota: 1000,
|
|
UsedQuota: 20,
|
|
RequestCount: 3,
|
|
AffCount: 2,
|
|
AffQuota: 800,
|
|
AffHistoryQuota: 1200,
|
|
}
|
|
user.SetAccessToken("old-token")
|
|
require.NoError(t, DB.Create(&user).Error)
|
|
|
|
staleUser, err := GetUserById(user.Id, true)
|
|
require.NoError(t, err)
|
|
|
|
require.NoError(t, DB.Model(&User{}).Where("id = ?", user.Id).Updates(map[string]interface{}{
|
|
"quota": gorm.Expr("quota - ?", 400),
|
|
"used_quota": gorm.Expr("used_quota + ?", 400),
|
|
"request_count": gorm.Expr("request_count + ?", 1),
|
|
"aff_count": gorm.Expr("aff_count + ?", 1),
|
|
"aff_quota": gorm.Expr("aff_quota - ?", 500),
|
|
"aff_history": gorm.Expr("aff_history + ?", 500),
|
|
"access_token": "rotated-token",
|
|
}).Error)
|
|
|
|
staleUser.DisplayName = "after"
|
|
require.NoError(t, staleUser.Update(false))
|
|
|
|
var got User
|
|
require.NoError(t, DB.First(&got, user.Id).Error)
|
|
assert.Equal(t, "after", got.DisplayName)
|
|
assert.Equal(t, 600, got.Quota)
|
|
assert.Equal(t, 420, got.UsedQuota)
|
|
assert.Equal(t, 4, got.RequestCount)
|
|
assert.Equal(t, 3, got.AffCount)
|
|
assert.Equal(t, 300, got.AffQuota)
|
|
assert.Equal(t, 1700, got.AffHistoryQuota)
|
|
assert.Equal(t, "rotated-token", got.GetAccessToken())
|
|
}
|
|
|
|
func TestUsageAccountingSupportsSignedDirectAndBatchDeltas(t *testing.T) {
|
|
setupUserUpdateTestState(t)
|
|
resetBatchUpdateTestState(t)
|
|
|
|
user := User{
|
|
Id: 10,
|
|
Username: "usage-adjustment-user",
|
|
Password: "password",
|
|
Status: common.UserStatusEnabled,
|
|
UsedQuota: 1000,
|
|
RequestCount: 3,
|
|
}
|
|
channel := Channel{
|
|
Id: 10,
|
|
Name: "usage-adjustment-channel",
|
|
Key: "sk-test",
|
|
Status: common.ChannelStatusEnabled,
|
|
UsedQuota: 1000,
|
|
}
|
|
require.NoError(t, DB.Create(&user).Error)
|
|
require.NoError(t, DB.Create(&channel).Error)
|
|
|
|
UpdateUserUsedQuota(user.Id, -200)
|
|
UpdateUserUsedQuota(user.Id, 50)
|
|
UpdateChannelUsedQuota(channel.Id, -200)
|
|
UpdateChannelUsedQuota(channel.Id, 50)
|
|
|
|
var got User
|
|
require.NoError(t, DB.Select("used_quota", "request_count").First(&got, user.Id).Error)
|
|
assert.Equal(t, 850, got.UsedQuota)
|
|
assert.Equal(t, 3, got.RequestCount)
|
|
var gotChannel Channel
|
|
require.NoError(t, DB.Select("used_quota").First(&gotChannel, channel.Id).Error)
|
|
assert.Equal(t, int64(850), gotChannel.UsedQuota)
|
|
|
|
common.BatchUpdateEnabled = true
|
|
UpdateUserUsedQuota(user.Id, 400)
|
|
UpdateUserUsedQuota(user.Id, -100)
|
|
UpdateChannelUsedQuota(channel.Id, 400)
|
|
UpdateChannelUsedQuota(channel.Id, -100)
|
|
|
|
require.NoError(t, DB.Select("used_quota", "request_count").First(&got, user.Id).Error)
|
|
assert.Equal(t, 850, got.UsedQuota, "batch deltas must remain queued until flush")
|
|
assert.Equal(t, 3, got.RequestCount)
|
|
require.NoError(t, DB.Select("used_quota").First(&gotChannel, channel.Id).Error)
|
|
assert.Equal(t, int64(850), gotChannel.UsedQuota, "batch deltas must remain queued until flush")
|
|
|
|
batchUpdate()
|
|
require.NoError(t, DB.Select("used_quota", "request_count").First(&got, user.Id).Error)
|
|
assert.Equal(t, 1150, got.UsedQuota)
|
|
assert.Equal(t, 3, got.RequestCount)
|
|
require.NoError(t, DB.Select("used_quota").First(&gotChannel, channel.Id).Error)
|
|
assert.Equal(t, int64(1150), gotChannel.UsedQuota)
|
|
}
|
|
|
|
func TestUpdateUserAccessTokenOnlyUpdatesAccessToken(t *testing.T) {
|
|
setupUserUpdateTestState(t)
|
|
|
|
user := User{
|
|
Id: 2,
|
|
Username: "token-rotation-user",
|
|
Password: "password",
|
|
DisplayName: "before",
|
|
Status: common.UserStatusEnabled,
|
|
Quota: 1000,
|
|
AffQuota: 800,
|
|
AffHistoryQuota: 1200,
|
|
}
|
|
require.NoError(t, DB.Create(&user).Error)
|
|
|
|
require.NoError(t, DB.Model(&User{}).Where("id = ?", user.Id).Updates(map[string]interface{}{
|
|
"quota": gorm.Expr("quota + ?", 500),
|
|
"aff_quota": gorm.Expr("aff_quota - ?", 500),
|
|
"display_name": "concurrent-update",
|
|
}).Error)
|
|
|
|
require.NoError(t, UpdateUserAccessToken(user.Id, "rotated-token"))
|
|
|
|
var got User
|
|
require.NoError(t, DB.First(&got, user.Id).Error)
|
|
assert.Equal(t, "rotated-token", got.GetAccessToken())
|
|
assert.Equal(t, "concurrent-update", got.DisplayName)
|
|
assert.Equal(t, 1500, got.Quota)
|
|
assert.Equal(t, 300, got.AffQuota)
|
|
assert.Equal(t, 1200, got.AffHistoryQuota)
|
|
}
|
|
|
|
func TestUpdateUserAccessTokenRejectsSoftDeletedUser(t *testing.T) {
|
|
setupUserUpdateTestState(t)
|
|
|
|
user := User{
|
|
Id: 3,
|
|
Username: "deleted-token-rotation-user",
|
|
Password: "password",
|
|
Status: common.UserStatusEnabled,
|
|
}
|
|
user.SetAccessToken("old-token")
|
|
require.NoError(t, DB.Create(&user).Error)
|
|
require.NoError(t, DB.Delete(&user).Error)
|
|
|
|
err := UpdateUserAccessToken(user.Id, "orphaned-token")
|
|
require.ErrorIs(t, err, gorm.ErrRecordNotFound)
|
|
|
|
var got User
|
|
require.NoError(t, DB.Unscoped().First(&got, user.Id).Error)
|
|
assert.Equal(t, "old-token", got.GetAccessToken())
|
|
}
|
|
|
|
func TestUpdateUserSettingOnlyUpdatesSetting(t *testing.T) {
|
|
setupUserUpdateTestState(t)
|
|
|
|
user := User{
|
|
Id: 2,
|
|
Username: "setting-user",
|
|
Password: "password",
|
|
Status: common.UserStatusEnabled,
|
|
Quota: 1000,
|
|
UsedQuota: 20,
|
|
RequestCount: 3,
|
|
}
|
|
require.NoError(t, DB.Create(&user).Error)
|
|
|
|
require.NoError(t, DB.Model(&User{}).Where("id = ?", user.Id).Updates(map[string]interface{}{
|
|
"quota": gorm.Expr("quota - ?", 250),
|
|
"used_quota": gorm.Expr("used_quota + ?", 250),
|
|
"request_count": gorm.Expr("request_count + ?", 1),
|
|
}).Error)
|
|
|
|
require.NoError(t, UpdateUserSetting(user.Id, dto.UserSetting{Language: "zh"}))
|
|
|
|
var got User
|
|
require.NoError(t, DB.First(&got, user.Id).Error)
|
|
assert.Equal(t, 750, got.Quota)
|
|
assert.Equal(t, 270, got.UsedQuota)
|
|
assert.Equal(t, 4, got.RequestCount)
|
|
assert.Equal(t, "zh", got.GetSetting().Language)
|
|
}
|
|
|
|
func TestEnsureEmailAvailableRejectsExistingEmailCaseInsensitive(t *testing.T) {
|
|
setupUserUpdateTestState(t)
|
|
|
|
require.NoError(t, DB.Create(&User{
|
|
Username: "existing",
|
|
Password: "old-password",
|
|
Email: "Taken@Example.com",
|
|
Status: common.UserStatusEnabled,
|
|
}).Error)
|
|
|
|
err := EnsureEmailAvailable(" taken@example.COM ", 0)
|
|
require.ErrorIs(t, err, ErrEmailAlreadyTaken)
|
|
|
|
user, err := GetUniqueUserByEmail("TAKEN@example.com")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "existing", user.Username)
|
|
|
|
require.NoError(t, EnsureEmailAvailable("taken@example.com", user.Id))
|
|
}
|
|
|
|
func TestInsertRejectsDuplicateEmailWithoutUniqueIndex(t *testing.T) {
|
|
setupUserUpdateTestState(t)
|
|
|
|
require.NoError(t, DB.Create(&User{
|
|
Username: "existing",
|
|
Password: "old-password",
|
|
Email: "taken@example.com",
|
|
Status: common.UserStatusEnabled,
|
|
}).Error)
|
|
|
|
user := &User{
|
|
Username: "oauth-user",
|
|
Email: "TAKEN@example.com",
|
|
Role: common.RoleCommonUser,
|
|
Status: common.UserStatusEnabled,
|
|
}
|
|
|
|
err := user.Insert(0)
|
|
require.ErrorIs(t, err, ErrEmailAlreadyTaken)
|
|
|
|
var count int64
|
|
require.NoError(t, DB.Model(&User{}).Where("username = ?", "oauth-user").Count(&count).Error)
|
|
assert.Zero(t, count)
|
|
}
|
|
|
|
func TestInsertKeepsBlankPasswordForPasswordlessUser(t *testing.T) {
|
|
setupUserUpdateTestState(t)
|
|
|
|
user := &User{
|
|
Username: "passwordless-user",
|
|
Role: common.RoleCommonUser,
|
|
Status: common.UserStatusEnabled,
|
|
}
|
|
|
|
require.NoError(t, user.Insert(0))
|
|
|
|
var stored User
|
|
require.NoError(t, DB.Where("username = ?", user.Username).First(&stored).Error)
|
|
assert.Empty(t, stored.Password)
|
|
}
|
|
|
|
func TestUpdateUserBindColumnOnlyTouchesTheBindingColumn(t *testing.T) {
|
|
truncateTables(t)
|
|
|
|
user := createUserBindTestUser(t)
|
|
require.NoError(t, DB.Model(&User{}).Where("id = ?", user.Id).Updates(map[string]interface{}{
|
|
"role": common.RoleAdminUser,
|
|
"status": common.UserStatusEnabled,
|
|
"group": "vip",
|
|
}).Error)
|
|
|
|
require.NoError(t, UpdateUserBindColumn(user.Id, "github_id", "gh-12345"))
|
|
|
|
reloaded, err := GetUserById(user.Id, true)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "gh-12345", reloaded.GitHubId)
|
|
assert.Equal(t, common.RoleAdminUser, reloaded.Role)
|
|
assert.Equal(t, common.UserStatusEnabled, reloaded.Status)
|
|
assert.Equal(t, "vip", reloaded.Group)
|
|
}
|
|
|
|
func TestUpdateUserBindColumnPreservesRestrictiveChange(t *testing.T) {
|
|
truncateTables(t)
|
|
|
|
user := createUserBindTestUser(t)
|
|
require.NoError(t, DB.Model(&User{}).Where("id = ?", user.Id).
|
|
Update("status", common.UserStatusDisabled).Error)
|
|
require.NoError(t, UpdateUserBindColumn(user.Id, "wechat_id", "wx-open-id"))
|
|
|
|
reloaded, err := GetUserById(user.Id, true)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "wx-open-id", reloaded.WeChatId)
|
|
assert.Equal(t, common.UserStatusDisabled, reloaded.Status)
|
|
}
|
|
|
|
func TestUpdateUserBindColumnRejectsNonWhitelistedColumns(t *testing.T) {
|
|
truncateTables(t)
|
|
|
|
user := createUserBindTestUser(t)
|
|
for _, column := range []string{"role", "status", "group", "quota", "username", "password", "id"} {
|
|
assert.Error(t, UpdateUserBindColumn(user.Id, column, "1"), "column %s must be rejected", column)
|
|
}
|
|
assert.Error(t, UpdateUserBindColumn(user.Id, "github_id; DROP TABLE users", "x"))
|
|
assert.Error(t, UpdateUserBindColumn(0, "github_id", "x"))
|
|
}
|
|
|
|
func TestValidateAndFillRejectsPasswordlessUser(t *testing.T) {
|
|
setupUserUpdateTestState(t)
|
|
|
|
require.NoError(t, DB.Create(&User{
|
|
Username: "passwordless-user",
|
|
Password: "",
|
|
Status: common.UserStatusEnabled,
|
|
}).Error)
|
|
|
|
loginUser := User{
|
|
Username: "passwordless-user",
|
|
Password: "NewPassword123",
|
|
}
|
|
err := loginUser.ValidateAndFill()
|
|
require.ErrorIs(t, err, ErrInvalidCredentials)
|
|
|
|
var stored User
|
|
require.NoError(t, DB.Where("username = ?", "passwordless-user").First(&stored).Error)
|
|
assert.Empty(t, stored.Password)
|
|
}
|
|
|
|
func TestResetUserPasswordByEmailRequiresSingleActiveMatch(t *testing.T) {
|
|
setupUserUpdateTestState(t)
|
|
|
|
require.NoError(t, DB.Create(&User{
|
|
Username: "duplicate-1",
|
|
Password: "old-1",
|
|
Email: "legacy@example.com",
|
|
AffCode: "dupe1",
|
|
Status: common.UserStatusEnabled,
|
|
}).Error)
|
|
require.NoError(t, DB.Create(&User{
|
|
Username: "duplicate-2",
|
|
Password: "old-2",
|
|
Email: "LEGACY@example.com",
|
|
AffCode: "dupe2",
|
|
Status: common.UserStatusEnabled,
|
|
}).Error)
|
|
|
|
err := ResetUserPasswordByEmail("legacy@example.com", "NewPassword123")
|
|
require.ErrorIs(t, err, ErrEmailAmbiguous)
|
|
|
|
var duplicates []User
|
|
require.NoError(t, DB.Where("LOWER(email) = ?", "legacy@example.com").Order("username asc").Find(&duplicates).Error)
|
|
require.Len(t, duplicates, 2)
|
|
assert.Equal(t, "old-1", duplicates[0].Password)
|
|
assert.Equal(t, "old-2", duplicates[1].Password)
|
|
|
|
require.NoError(t, DB.Create(&User{
|
|
Username: "unique",
|
|
Password: "old",
|
|
Email: "unique@example.com",
|
|
AffCode: "unique",
|
|
Status: common.UserStatusEnabled,
|
|
}).Error)
|
|
|
|
require.NoError(t, ResetUserPasswordByEmail("UNIQUE@example.com", "NewPassword123"))
|
|
|
|
var unique User
|
|
require.NoError(t, DB.Where("username = ?", "unique").First(&unique).Error)
|
|
assert.True(t, common.ValidatePasswordAndHash("NewPassword123", unique.Password))
|
|
|
|
err = ResetUserPasswordByEmail("missing@example.com", "NewPassword123")
|
|
require.True(t, errors.Is(err, ErrEmailNotFound))
|
|
}
|