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) }