sub2api/backend/internal/repository/account_traffic_cache_test.go

package repository

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/Wei-Shaw/sub2api/internal/service"
	"github.com/alicebob/miniredis/v2"
	"github.com/redis/go-redis/v9"
	"github.com/stretchr/testify/require"
)

func trafficFixture(t *testing.T) (*miniredis.Miniredis, *accountTrafficCache, service.AccountTrafficPlan) {
	t.Helper()
	server := miniredis.RunT(t)
	server.SetTime(time.Date(2026, 9, 14, 12, 0, 0, 0, time.UTC))
	client := redis.NewClient(&redis.Options{Addr: server.Addr()})
	t.Cleanup(func() { client.Close() })
	p := service.DefaultAccountTrafficPolicy()
	p.StrictRPMEnabled = true
	p.RPM = 3
	p.Burst = 3
	return server, &accountTrafficCache{rdb: client}, service.AccountTrafficPlan{AccountID: 42, Revision: 1, Signature: "policy-1", HardLimit: 8, Policy: p}
}

func TestAccountTrafficAtomicRPMAndBurst(t *testing.T) {
	server, cache, plan := trafficFixture(t)
	ctx := context.Background()
	for i := 0; i < 3; i++ {
		a, err := cache.Acquire(ctx, plan, fmt.Sprint(i))
		require.NoError(t, err)
		require.True(t, a.Allowed)
		require.NoError(t, cache.Finish(ctx, plan, fmt.Sprint(i), 200, 100))
	}
	blocked, err := cache.Acquire(ctx, plan, "fourth")
	require.NoError(t, err)
	require.False(t, blocked.Allowed)
	require.Equal(t, time.Minute, blocked.RetryAfter)
	server.SetTime(time.Date(2026, 9, 14, 12, 0, 20, 0, time.UTC))
	blocked, err = cache.Acquire(ctx, plan, "still-full")
	require.NoError(t, err)
	require.False(t, blocked.Allowed)
	require.Equal(t, 40*time.Second, blocked.RetryAfter)
	server.SetTime(time.Date(2026, 9, 14, 12, 1, 1, 0, time.UTC))
	accepted, err := cache.Acquire(ctx, plan, "new-window")
	require.NoError(t, err)
	require.True(t, accepted.Allowed)
	state, err := cache.Snapshot(ctx, plan)
	require.NoError(t, err)
	require.Equal(t, 1, state.RequestsLastMinute)
	require.EqualValues(t, 4, state.Accepted)
	// A burst smaller than RPM stops a simultaneous spike before the minute cap.
	plan.AccountID = 43
	plan.Policy.RPM = 60
	plan.Policy.Burst = 2
	for i := 0; i < 2; i++ {
		a, err := cache.Acquire(ctx, plan, fmt.Sprint(i))
		require.NoError(t, err)
		require.True(t, a.Allowed)
	}
	blocked, err = cache.Acquire(ctx, plan, "burst-full")
	require.NoError(t, err)
	require.False(t, blocked.Allowed)
	require.Equal(t, time.Second, blocked.RetryAfter)
}

func TestAccountTrafficConcurrentInstancesCannotExceedCeiling(t *testing.T) {
	server, cache, plan := trafficFixture(t)
	plan.Policy.RPM = 1
	plan.Policy.Burst = 1
	otherClient := redis.NewClient(&redis.Options{Addr: server.Addr()})
	defer otherClient.Close()
	other := &accountTrafficCache{rdb: otherClient}
	var admitted atomic.Int32
	var wg sync.WaitGroup
	failures := make(chan error, 40)
	for i := 0; i < 40; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			chosen := cache
			if i%2 == 1 {
				chosen = other
			}
			a, err := chosen.Acquire(context.Background(), plan, fmt.Sprint(i))
			if err != nil {
				failures <- err
			}
			if a.Allowed {
				admitted.Add(1)
			}
		}(i)
	}
	wg.Wait()
	close(failures)
	for err := range failures {
		require.NoError(t, err)
	}
	require.EqualValues(t, 1, admitted.Load())
}

func TestAccountTrafficAdaptiveSuggestionAndSlowRecovery(t *testing.T) {
	server, cache, plan := trafficFixture(t)
	ctx := context.Background()
	plan.Policy.StrictRPMEnabled = false
	plan.Policy.AdaptiveEnabled = true
	plan.Policy.FailureThreshold = 2
	for i := 0; i < 2; i++ {
		id := fmt.Sprint(i)
		a, err := cache.Acquire(ctx, plan, id)
		require.NoError(t, err)
		require.True(t, a.Allowed)
		require.NoError(t, cache.Finish(ctx, plan, id, 429, 100))
	}
	state, err := cache.Snapshot(ctx, plan)
	require.NoError(t, err)
	require.Equal(t, 4, state.RecommendedConcurrency)
	require.Equal(t, 8, state.EffectiveConcurrency, "observe mode never lowers the actual cap")
	plan.Revision = 2
	plan.Signature = "automatic"
	plan.Policy.AdaptiveMode = "automatic"
	require.NoError(t, cache.Sync(ctx, plan))
	for i := 0; i < 4; i++ {
		a, err := cache.Acquire(ctx, plan, fmt.Sprint("active-", i))
		require.NoError(t, err)
		require.True(t, a.Allowed)
	}
	blocked, err := cache.Acquire(ctx, plan, "fifth")
	require.NoError(t, err)
	require.False(t, blocked.Allowed)
	require.Equal(t, 4, blocked.State.EffectiveConcurrency)
	for i := 0; i < 3; i++ {
		require.NoError(t, cache.Finish(ctx, plan, fmt.Sprint("active-", i), 200, 100))
	}
	state, err = cache.Snapshot(ctx, plan)
	require.NoError(t, err)
	require.Equal(t, 4, state.EffectiveConcurrency)
	server.SetTime(time.Date(2026, 9, 14, 12, 1, 1, 0, time.UTC))
	require.NoError(t, cache.Finish(ctx, plan, "active-3", 200, 100))
	state, err = cache.Snapshot(ctx, plan)
	require.NoError(t, err)
	require.Equal(t, 5, state.EffectiveConcurrency)
	require.Equal(t, 8, plan.HardLimit, "the saved administrator cap is not mutated")
	plan.Policy.AdaptiveEnabled = false
	plan.Revision = 3
	plan.Signature = "disabled"
	require.NoError(t, cache.Sync(ctx, plan))
	state, err = cache.Snapshot(ctx, plan)
	require.NoError(t, err)
	require.Equal(t, 8, state.EffectiveConcurrency)
}

func TestAccountTrafficStalePolicyAndOldCompletionCannotUndoChanges(t *testing.T) {
	_, cache, old := trafficFixture(t)
	ctx := context.Background()
	old.Policy.StrictRPMEnabled = false
	old.Policy.AdaptiveEnabled = true
	old.Policy.FailureThreshold = 1
	a, err := cache.Acquire(ctx, old, "old-request")
	require.NoError(t, err)
	require.True(t, a.Allowed)
	updated := old
	updated.Revision = 2
	updated.Signature = "turned-off"
	updated.Policy.AdaptiveEnabled = false
	require.NoError(t, cache.Sync(ctx, updated))
	blocked, err := cache.Acquire(ctx, old, "stale")
	require.NoError(t, err)
	require.True(t, blocked.Allowed)
	require.True(t, blocked.SkipObservation, "turning off observation must not block an old session")
	require.NoError(t, cache.Finish(ctx, old, "old-request", 429, 50))
	state, err := cache.Snapshot(ctx, updated)
	require.NoError(t, err)
	require.Equal(t, 8, state.EffectiveConcurrency)
	require.Equal(t, 8, state.RecommendedConcurrency)
	require.Zero(t, state.InFlight)
}