service/notify_limit_test.go

Rama main · réplica de lectura del repositorio original.

package service

import (
	"fmt"
	"testing"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/relaykit/dto"
	"github.com/alicebob/miniredis/v2"
	"github.com/go-redis/redis/v8"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// 限流窗口必须由 NOTIFICATION_LIMIT_DURATION_MINUTE 决定，不能在整点被重置
func TestCheckNotificationLimitWindowSpansClockHours(t *testing.T) {
	originalRedis := common.RedisEnabled
	originalCount := constant.NotifyLimitCount
	originalMinute := constant.NotificationLimitDurationMinute
	t.Cleanup(func() {
		common.RedisEnabled = originalRedis
		constant.NotifyLimitCount = originalCount
		constant.NotificationLimitDurationMinute = originalMinute
	})
	common.RedisEnabled = false
	constant.NotifyLimitCount = 2
	constant.NotificationLimitDurationMinute = 720

	tests := []struct {
		name        string
		windowStart time.Duration
		expected    bool
	}{
		{name: "limit reached two hours ago still blocks", windowStart: 2 * time.Hour, expected: false},
		{name: "limit reached after window expired allows", windowStart: 721 * time.Minute, expected: true},
	}
	for i, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			userId := 900001 + i
			key := fmt.Sprintf("%d:%s", userId, dto.NotifyTypeQuotaExceed)
			notifyLimitStore.Store(key, limitCount{Count: 2, Timestamp: time.Now().Add(-tt.windowStart)})
			t.Cleanup(func() { notifyLimitStore.Delete(key) })

			canSend, err := CheckNotificationLimit(userId, dto.NotifyTypeQuotaExceed)
			require.NoError(t, err)
			assert.Equal(t, tt.expected, canSend)
		})
	}
}

func TestCheckNotificationLimitRedisWindowSpansClockHours(t *testing.T) {
	originalRedis := common.RedisEnabled
	originalRDB := common.RDB
	originalCount := constant.NotifyLimitCount
	originalMinute := constant.NotificationLimitDurationMinute
	server := miniredis.RunT(t)
	common.RDB = redis.NewClient(&redis.Options{Addr: server.Addr(), MaxRetries: -1})
	t.Cleanup(func() {
		_ = common.RDB.Close()
		common.RDB = originalRDB
		common.RedisEnabled = originalRedis
		constant.NotifyLimitCount = originalCount
		constant.NotificationLimitDurationMinute = originalMinute
	})
	common.RedisEnabled = true
	constant.NotifyLimitCount = 2
	constant.NotificationLimitDurationMinute = 720

	// 窗口在 2 小时前开始并已发满，剩余 10 小时
	userId := 900010
	key := fmt.Sprintf("notify_limit:%d:%s", userId, dto.NotifyTypeQuotaExceed)
	require.NoError(t, server.Set(key, "2"))
	server.SetTTL(key, 10*time.Hour)

	steps := []struct {
		name     string
		elapse   time.Duration
		expected bool
	}{
		{name: "limit reached two hours ago still blocks", expected: false},
		{name: "new window opens after the old one ends", elapse: 10*time.Hour + time.Minute, expected: true},
		{name: "second notification in the new window", expected: true},
		{name: "third notification in the new window is blocked", expected: false},
		{name: "new window ends after the configured duration", elapse: 720 * time.Minute, expected: true},
	}
	for _, step := range steps {
		server.FastForward(step.elapse)
		canSend, err := CheckNotificationLimit(userId, dto.NotifyTypeQuotaExceed)
		require.NoError(t, err, step.name)
		assert.Equal(t, step.expected, canSend, step.name)
	}
}