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