93 lines
2.5 KiB
Go
93 lines
2.5 KiB
Go
package model
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/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 TestUserUpdateDoesNotOverwriteAccountingFields(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,
|
|
}
|
|
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),
|
|
}).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)
|
|
}
|
|
|
|
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)
|
|
}
|