service/user_notify_test.go
Rama main · réplica de lectura del repositorio original.
package service
import (
"fmt"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
// A low-balance reminder goes out once, on the charge that takes a wallet or
// subscription balance below the threshold. A balance already below it earns
// nothing more until a top-up lifts it back above the line. The site switch
// and the user's switch both turn reminders off; a user who never chose gets
// them.
func TestQuotaReminderFiresOncePerDrop(t *testing.T) {
previousEnabled, previousThreshold := common.QuotaRemindEnabled, common.QuotaRemindThreshold
t.Cleanup(func() {
common.QuotaRemindEnabled, common.QuotaRemindThreshold = previousEnabled, previousThreshold
})
common.QuotaRemindThreshold = 1000
on, off := true, false
// changes move the balance from start: negative ones are charges, which
// may remind; positive ones are top-ups. want holds one entry per charge.
for _, tc := range []struct {
name string
siteEnabled bool
setting dto.UserSetting
start int64
changes []int64
want []bool
}{
{
name: "never chosen: one reminder per drop, again after a top-up",
siteEnabled: true,
start: 1500,
changes: []int64{-400, -200, -300, +1500, -1200},
want: []bool{false, true, false, true},
},
{
name: "switched on with its own threshold",
siteEnabled: true,
setting: dto.UserSetting{QuotaWarningEnabled: &on, QuotaWarningThreshold: 5000},
start: 6000,
changes: []int64{-500, -600, -2000},
want: []bool{false, true, false},
},
{
name: "a charge that lands exactly on the threshold does not remind",
siteEnabled: true,
start: 1200,
changes: []int64{-200, -1},
want: []bool{false, true},
},
{
name: "user switched off",
siteEnabled: true,
setting: dto.UserSetting{QuotaWarningEnabled: &off},
start: 1500,
changes: []int64{-600},
want: []bool{false},
},
{
name: "site switched off",
setting: dto.UserSetting{QuotaWarningEnabled: &on},
start: 1500,
changes: []int64{-600},
want: []bool{false},
},
} {
t.Run(tc.name, func(t *testing.T) {
common.QuotaRemindEnabled = tc.siteEnabled
got := make([]bool, 0, len(tc.want))
balance := tc.start
for _, change := range tc.changes {
if change < 0 {
got = append(got, quotaReminderDue(tc.setting, balance, balance+change))
}
balance += change
}
assert.Equal(t, tc.want, got)
})
}
}
// Channel notices go to every enabled administrator whose channel switch is
// on, in that administrator's saved language. Until an administrator chooses,
// root receives them and other administrators do not; other users never do.
func TestChannelNoticesFollowEachAdministratorsSwitch(t *testing.T) {
previousDB := model.DB
previousCache, previousRedis := common.MemoryCacheEnabled, common.RedisEnabled
previousNotifyLimit := constant.NotifyLimitCount
previousClient, previousWorker := httpClient, system_setting.WorkerUrl
fetch := system_setting.GetFetchSetting()
previousFetch := *fetch
t.Cleanup(func() {
model.DB = previousDB
common.MemoryCacheEnabled, common.RedisEnabled = previousCache, previousRedis
constant.NotifyLimitCount = previousNotifyLimit
httpClient, system_setting.WorkerUrl = previousClient, previousWorker
*fetch = previousFetch
})
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
sqlDB, err := database.DB()
require.NoError(t, err)
sqlDB.SetMaxOpenConns(1)
t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
require.NoError(t, database.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.User{}))
model.DB = database
common.MemoryCacheEnabled, common.RedisEnabled = false, false
constant.NotifyLimitCount = 100
var mutex sync.Mutex
received := map[string][]string{}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var payload WebhookPayload
if err := common.DecodeJson(r.Body, &payload); err == nil {
mutex.Lock()
recipient := strings.TrimPrefix(r.URL.Path, "/")
received[recipient] = append(received[recipient], payload.Type+" "+payload.Title)
mutex.Unlock()
}
w.WriteHeader(http.StatusNoContent)
}))
t.Cleanup(server.Close)
httpClient, system_setting.WorkerUrl = server.Client(), ""
fetch.EnableSSRFProtection = false
on, off := true, false
for _, user := range []struct {
name string
role int
status int
receives *bool
language string
}{
{"root-never-chose", common.RoleRootUser, common.UserStatusEnabled, nil, ""},
{"root-turned-off", common.RoleRootUser, common.UserStatusEnabled, &off, ""},
{"admin-never-chose", common.RoleAdminUser, common.UserStatusEnabled, nil, ""},
{"admin-turned-on", common.RoleAdminUser, common.UserStatusEnabled, &on, "zhCN"},
{"disabled-admin-turned-on", common.RoleAdminUser, common.UserStatusDisabled, &on, ""},
{"user-with-switch-stored-on", common.RoleCommonUser, common.UserStatusEnabled, &on, ""},
} {
setting, err := common.Marshal(dto.UserSetting{
NotifyType: dto.NotifyTypeWebhook,
WebhookUrl: server.URL + "/" + user.name,
ChannelNotifyEnabled: user.receives,
Language: user.language,
})
require.NoError(t, err)
require.NoError(t, database.Create(&model.User{Username: user.name, Role: user.role, Status: user.status, Setting: string(setting), AffCode: user.name}).Error)
}
channel := &model.Channel{Name: "notice-review", Key: "fixture-key", Type: 1, Status: common.ChannelStatusEnabled, Group: "default", Models: "test-model"}
require.NoError(t, channel.Insert())
DisableChannel(types.ChannelError{ChannelId: channel.Id, ChannelName: channel.Name, AutoBan: true}, "upstream failed")
EnableChannel(channel.Id, "", channel.Name)
NotifyChannelWatchers(dto.NotifyTypeChannelTest, func(lang string) (string, string) {
return i18n.Translate(lang, i18n.MsgChannelTestCompletedSubject), i18n.Translate(lang, i18n.MsgChannelTestCompletedContent)
})
disabledType := formatNotifyType(channel.Id, common.ChannelStatusAutoDisabled)
enabledType := formatNotifyType(channel.Id, common.ChannelStatusEnabled)
assert.Equal(t, map[string][]string{
"root-never-chose": {
fmt.Sprintf(`%s Channel "notice-review" (#%d) has been disabled`, disabledType, channel.Id),
fmt.Sprintf(`%s Channel "notice-review" (#%d) has been enabled`, enabledType, channel.Id),
"channel_test Channel test completed",
},
"admin-turned-on": {
fmt.Sprintf("%s 通道「notice-review」（#%d）已被禁用", disabledType, channel.Id),
fmt.Sprintf("%s 通道「notice-review」（#%d）已被启用", enabledType, channel.Id),
"channel_test 通道测试完成",
},
}, received)
}