nezha_domains/cmd/dashboard/controller/mcp_ratelimit_boundary_test.go

package controller

import (
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

type mcpRateLimitFakeClock struct {
	now time.Time
}

func (clock *mcpRateLimitFakeClock) Now() time.Time {
	return clock.now
}

func (clock *mcpRateLimitFakeClock) Advance(duration time.Duration) {
	clock.now = clock.now.Add(duration)
}

func TestMCPRateLimiter_baselinePreservesAnonymousBypassAndBucketReset(t *testing.T) {
	// Given
	limiter := newMCPRateLimiter(3, 3)

	// When
	anonymousAllowed := limiter.Allow(0)
	firstAllowed := limiter.Allow(7)
	secondAllowed := limiter.Allow(7)
	thirdAllowed := limiter.Allow(7)
	fourthRejected := limiter.Allow(7)
	// Then
	if !anonymousAllowed || !firstAllowed || !secondAllowed || !thirdAllowed || fourthRejected {
		t.Fatalf("baseline limiter semantics changed: anonymous=%t first=%t second=%t third=%t rejected=%t", anonymousAllowed, firstAllowed, secondAllowed, thirdAllowed, fourthRejected)
	}
}

func TestMCPRateLimiter_allowsTenAndRejectsElevenWithinOneSecond(t *testing.T) {
	// Given
	clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)}
	limiter := newMCPRateLimiterWithClock(10, 120, clock.Now)

	// When
	allowed := 0
	for requestNumber := 1; requestNumber <= 11; requestNumber++ {
		if limiter.Allow(7) {
			allowed++
		}
	}

	// Then
	if allowed != 10 {
		t.Fatalf("allowed count = %d, want 10", allowed)
	}
	if limiter.Allow(7) {
		t.Fatal("twelfth request must remain rejected in the same second")
	}
}

func TestMCPRateLimiter_allowsOneAfterSecondBucketRollover(t *testing.T) {
	// Given
	clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)}
	limiter := newMCPRateLimiterWithClock(10, 120, clock.Now)
	for requestNumber := 1; requestNumber <= 10; requestNumber++ {
		if !limiter.Allow(7) {
			t.Fatalf("request %d must be allowed before rollover", requestNumber)
		}
	}

	// When
	clock.Advance(time.Second)

	// Then
	if !limiter.Allow(7) {
		t.Fatal("request after one-second bucket rollover must be allowed")
	}
}

func TestMCPRateLimiter_allows120AndRejects121WithinOneMinute(t *testing.T) {
	// Given
	clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)}
	limiter := newMCPRateLimiterWithClock(10_000, 120, clock.Now)

	// When
	allowed := 0
	for requestNumber := 1; requestNumber <= 121; requestNumber++ {
		if limiter.Allow(7) {
			allowed++
		}
		clock.Advance(100 * time.Millisecond)
	}

	// Then
	if allowed != 120 {
		t.Fatalf("allowed count = %d, want 120", allowed)
	}
}

func TestMCPRateLimiter_independentTokensKeepIndependentBudgets(t *testing.T) {
	// Given
	clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)}
	limiter := newMCPRateLimiterWithClock(1, 120, clock.Now)

	// When
	firstTokenAllowed := limiter.Allow(7)
	firstTokenRejected := limiter.Allow(7)
	secondTokenAllowed := limiter.Allow(8)

	// Then
	if !firstTokenAllowed || firstTokenRejected || !secondTokenAllowed {
		t.Fatalf("token budgets are not independent: first allowed=%t rejected=%t second allowed=%t", firstTokenAllowed, firstTokenRejected, secondTokenAllowed)
	}
}

func TestMCPRateLimiter_clockControlledCallsDoNotMutateSharedLimiter(t *testing.T) {
	// Given
	originalLimiter := mcpRateLimiterShared
	clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)}
	limiter := newMCPRateLimiterWithClock(1, 1, clock.Now)

	// When
	if !limiter.Allow(7) {
		t.Fatal("isolated limiter request must be allowed")
	}

	// Then
	if mcpRateLimiterShared != originalLimiter {
		t.Fatal("isolated limiter changed shared production limiter ownership")
	}
}

func TestMCPRateLimiter_concurrentCallsAtSecondBoundaryAllowExactlyLimit(t *testing.T) {
	// Given
	clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)}
	limiter := newMCPRateLimiterWithClock(10, 120, clock.Now)
	var allowed atomic.Int32
	var waitGroup sync.WaitGroup

	// When
	for requestNumber := 0; requestNumber < 20; requestNumber++ {
		waitGroup.Add(1)
		go func() {
			defer waitGroup.Done()
			if limiter.Allow(7) {
				allowed.Add(1)
			}
		}()
	}
	waitGroup.Wait()

	// Then
	if got := allowed.Load(); got != 10 {
		t.Fatalf("concurrent allowed count = %d, want 10", got)
	}
}

func TestMCPRateLimiter_minuteBucketRolloverRestoresBudget(t *testing.T) {
	// Given
	clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)}
	limiter := newMCPRateLimiterWithClock(10_000, 1, clock.Now)
	if !limiter.Allow(7) {
		t.Fatal("first request must be allowed")
	}
	if limiter.Allow(7) {
		t.Fatal("second request must be rejected before minute rollover")
	}

	// When
	clock.Advance(time.Minute)

	// Then
	if !limiter.Allow(7) {
		t.Fatal("request after minute rollover must be allowed")
	}
}

func TestMCPRateLimiter_samplesClockWhileHoldingStateLock(t *testing.T) {
	// Given
	now := time.Unix(1_700_000_000, 0)
	var limiter *MCPRateLimiter
	clock := func() time.Time {
		if limiter.mu.TryLock() {
			limiter.mu.Unlock()
			t.Fatal("clock callback acquired limiter state lock; Allow sampled before locking")
		}
		return now
	}
	limiter = newMCPRateLimiterWithClock(1, 1, clock)

	// When
	allowed := limiter.Allow(7)

	// Then
	if !allowed {
		t.Fatal("initial request must be allowed")
	}
}