sub2api/backend/internal/repository/scheduler_threshold_cache_unit_test.go

//go:build unit

package repository

import (
	"context"
	"testing"
	"time"

	"github.com/Wei-Shaw/sub2api/internal/config"
	"github.com/Wei-Shaw/sub2api/internal/service"
	"github.com/stretchr/testify/require"
)

// Only persistence is stubbed: the candidate is serialized through the real
// Redis snapshot projection before the production threshold service sees it.
type snapshotThresholdRepo struct {
	service.AccountRepository
	accountPauses int
	modelPauses   int
	until         time.Time
}

func (r *snapshotThresholdRepo) SetTempUnschedulable(_ context.Context, _ int64, until time.Time, _ string) error {
	r.accountPauses++
	r.until = until
	return nil
}

func (r *snapshotThresholdRepo) SetModelRateLimit(_ context.Context, _ int64, _ string, until time.Time, _ ...string) error {
	r.modelPauses++
	r.until = until
	return nil
}

func TestSchedulerCacheAnthropicThresholdAdmission(t *testing.T) {
	now := time.Now().UTC().Truncate(time.Second)
	reset := now.Add(time.Hour)
	cases := []struct {
		name                       string
		fiveHour, sevenDay, fable  float64
		threshold                  int
		expired                    bool
		accountPaused, fablePaused bool
	}{
		{name: "shared 5h", fiveHour: .60, threshold: 60, accountPaused: true},
		{name: "shared 7d", sevenDay: .66, threshold: 60, accountPaused: true},
		{name: "Fable only", sevenDay: .40, fable: .75, threshold: 60, fablePaused: true},
		{name: "below threshold", fiveHour: .59, sevenDay: .59, fable: .59, threshold: 60},
		{name: "expired windows", fiveHour: .75, sevenDay: .75, fable: .75, threshold: 60, expired: true},
		{name: "disabled override", fiveHour: .75, sevenDay: .75, fable: .75, threshold: 100},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			ctx := context.Background()
			end := reset
			if tc.expired {
				end = now.Add(-time.Hour)
			}
			account := service.Account{
				ID: 3, Platform: service.PlatformAnthropic, Type: service.AccountTypeOAuth,
				Status: service.StatusActive, Schedulable: true, SessionWindowEnd: &end,
				Credentials: map[string]any{"account_scheduling_threshold": tc.threshold, "access_token": "test-only-token"},
				Extra: map[string]any{
					"session_window_utilization":   tc.fiveHour,
					"passive_usage_7d_utilization": tc.sevenDay, "passive_usage_7d_reset": end.Unix(),
					"passive_usage_7d_oi_utilization": tc.fable, "passive_usage_7d_oi_reset": end.Unix(),
					"unrelated_large_payload": "drop me",
				},
			}
			cache := newSchedulerCacheUnit(t)
			bucket := service.SchedulerBucket{GroupID: 8, Platform: service.PlatformAnthropic, Mode: service.SchedulerModeSingle}
			token, err := cache.CaptureBucketWriteToken(ctx, bucket)
			require.NoError(t, err)
			require.NoError(t, cache.SetSnapshot(ctx, bucket, token, []service.Account{account}))
			candidates, hit, err := cache.GetSnapshot(ctx, bucket)
			require.NoError(t, err)
			require.True(t, hit)
			require.Len(t, candidates, 1)
			require.NotContains(t, candidates[0].Credentials, "access_token")
			require.NotContains(t, candidates[0].Extra, "unrelated_large_payload")
			full, err := cache.GetAccount(ctx, account.ID)
			require.NoError(t, err)
			require.NotNil(t, full)
			for _, candidate := range []*service.Account{candidates[0], full} {
				repo := &snapshotThresholdRepo{}
				limiter := service.NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
				// Default platform thresholds are disabled, so the per-account override
				// must survive both cache read paths to trigger admission control.
				limiter.SetSettingService(service.NewSettingService(nil, &config.Config{}))
				require.Equal(t, tc.accountPaused, limiter.ApplyAccountSchedulingThreshold(ctx, candidate))
				require.Equal(t, tc.accountPaused, repo.accountPauses == 1)
				require.Equal(t, tc.fablePaused, repo.modelPauses == 1)
				if tc.accountPaused || tc.fablePaused {
					require.True(t, end.Equal(repo.until))
				}
				if tc.fablePaused {
					require.False(t, candidate.IsSchedulableForModel("claude-fable-5"))
					require.True(t, candidate.IsSchedulableForModel("claude-opus-5"))
				}
			}
		})
	}
}

func TestSchedulerCacheAnthropicUsageRefresh(t *testing.T) {
	ctx := context.Background()
	now := time.Now().UTC().Truncate(time.Second)
	end := now.Add(time.Hour)
	cache := newSchedulerCacheUnit(t)
	bucket := service.SchedulerBucket{GroupID: 8, Platform: service.PlatformAnthropic, Mode: service.SchedulerModeSingle}
	account := service.Account{
		ID: 3, Platform: service.PlatformAnthropic, Type: service.AccountTypeOAuth,
		Status: service.StatusActive, Schedulable: true,
		Credentials: map[string]any{"account_scheduling_threshold": 60},
		Extra:       map[string]any{"passive_usage_7d_utilization": .59, "passive_usage_7d_reset": end.Unix()},
	}
	token, err := cache.CaptureBucketWriteToken(ctx, bucket)
	require.NoError(t, err)
	require.NoError(t, cache.SetSnapshot(ctx, bucket, token, []service.Account{account}))
	for _, used := range []float64{.59, .66, .10} {
		account.Extra["passive_usage_7d_utilization"] = used
		// UpdateExtra refreshes account payloads without rebuilding bucket membership.
		require.NoError(t, cache.SetAccount(ctx, &account))
		candidates, hit, err := cache.GetSnapshot(ctx, bucket)
		require.NoError(t, err)
		require.True(t, hit)
		require.Len(t, candidates, 1)
		decision := service.EvaluateAccountSchedulingThreshold(candidates[0], map[string]int{service.PlatformAnthropic: 100}, now)
		require.Equal(t, used >= .60, decision.ShouldPause)
	}
}