package youtube import ( "context" "sync" "sync/atomic" "testing" "time" "github.com/stretchr/testify/require" "golang.org/x/time/rate" ) // waitOn drives a test-scoped limiter the same way WaitFetchGate drives the // global one, so these tests exercise the gate's behaviour without mutating the // process-wide gate (which would pollute sibling tests / the click-path). func waitOn(t *testing.T, lim *rate.Limiter) func(context.Context) error { t.Helper() return func(ctx context.Context) error { return lim.Wait(ctx) } } func TestFetchGateSerialisesConcurrentCallers(t *testing.T) { const ( n = 5 interval = 10 * time.Millisecond ) lim := rate.NewLimiter(rate.Every(interval), 1) wait := waitOn(t, lim) var ( inFlight, maxInFlight atomic.Int32 wg sync.WaitGroup ) start := time.Now() for i := 0; i < n; i++ { wg.Add(1) go func() { defer wg.Done() require.NoError(t, wait(context.Background())) cur := inFlight.Add(1) for { old := maxInFlight.Load() if cur <= old || maxInFlight.CompareAndSwap(old, cur) { break } } // Hold the "critical section" briefly so overlap would be observable. time.Sleep(interval / 4) inFlight.Add(-1) }() } wg.Wait() elapsed := time.Since(start) require.Equal(t, int32(1), maxInFlight.Load(), "the gate must admit at most one caller per interval — no overlap") require.GreaterOrEqual(t, elapsed, time.Duration(n-1)*interval, "N gated callers take at least (N-1)*interval wall time") } func TestFetchGateRespectsContextCancellation(t *testing.T) { // A slow gate (1 token/hour, burst already spent) blocks; a cancelled ctx must // unblock Wait with an error rather than hang. lim := rate.NewLimiter(rate.Every(time.Hour), 1) require.True(t, lim.Allow(), "spend the single burst token") ctx, cancel := context.WithCancel(context.Background()) cancel() require.Error(t, waitOn(t, lim)(ctx), "cancelled ctx must fail Wait, not block") } func TestSetFetchRateZeroIsUnlimited(t *testing.T) { // Snapshot and restore the global so this test does not pollute the process. prev := globalFetchGate t.Cleanup(func() { globalFetchGate = prev }) SetFetchRate(0) require.Equal(t, rate.Inf, globalFetchGate.Limit(), "0 interval = unlimited gate") // An unlimited gate never blocks, even back-to-back. for i := 0; i < 100; i++ { require.NoError(t, WaitFetchGate(context.Background())) } } func TestForegroundContextMarker(t *testing.T) { require.False(t, isForeground(context.Background()), "plain context is background") require.True(t, isForeground(ForegroundContext(context.Background())), "marked context is foreground") } // TestWaitFetchGateForegroundProceedsImmediately: a foreground fetch acquires a // token without yielding, even when background callers exist. func TestWaitFetchGateForegroundProceedsImmediately(t *testing.T) { SetFetchRate(0) // unlimited limiter — isolate the yield logic from pacing foregroundPending.Store(0) t.Cleanup(func() { foregroundPending.Store(0) }) ctx, cancel := context.WithTimeout(context.Background(), time.Second) defer cancel() require.NoError(t, WaitFetchGate(ForegroundContext(ctx)), "foreground proceeds immediately") } // TestWaitFetchGateBackgroundYieldsToForeground: while a foreground fetch is // pending, a background fetch yields (does not take a token) until the foreground // clears — proven by a background wait timing out against its own deadline, then // succeeding once the foreground is done. func TestWaitFetchGateBackgroundYieldsToForeground(t *testing.T) { SetFetchRate(0) foregroundPending.Store(1) // simulate a foreground fetch in flight t.Cleanup(func() { foregroundPending.Store(0) }) ctx, cancel := context.WithTimeout(context.Background(), 250*time.Millisecond) defer cancel() require.Error(t, WaitFetchGate(ctx), "background yields (blocks) while foreground is pending") foregroundPending.Store(0) // foreground done require.NoError(t, WaitFetchGate(context.Background()), "background proceeds once foreground clears") }