TREK/server/tests/unit/plugins/rate-limit.test.ts

/**
 * Per-plugin RPC rate limiter (#plugins hardening): a token bucket + an in-flight
 * gauge, clock injected so it tests without timers.
 */
import { TokenBucket, RpcRateLimiter } from '../../../src/nest/plugins/host/rate-limit';

import { describe, it, expect } from 'vitest';

describe('TokenBucket', () => {
  it('allows a burst up to capacity, then throttles until it refills', () => {
    const b = new TokenBucket(3, 1, 0); // 3 burst, 1/sec
    expect(b.take(0)).toBe(true);
    expect(b.take(0)).toBe(true);
    expect(b.take(0)).toBe(true);
    expect(b.take(0)).toBe(false); // burst spent
    // 1s later one token has refilled
    expect(b.take(1000)).toBe(true);
    expect(b.take(1000)).toBe(false);
  });

  it('never refills above capacity even after a long idle', () => {
    const b = new TokenBucket(2, 5, 0);
    // idle 10s would add 50 tokens uncapped — must cap at 2
    expect(b.take(10_000)).toBe(true);
    expect(b.take(10_000)).toBe(true);
    expect(b.take(10_000)).toBe(false);
  });
});

describe('RpcRateLimiter', () => {
  it('rate-limits via the bucket and reserves/releases in-flight slots', () => {
    const l = new RpcRateLimiter({ burst: 2, perSec: 1, maxInFlight: 5 }, 0);
    expect(l.tryAcquire(0)).toBe(true);
    expect(l.tryAcquire(0)).toBe(true);
    expect(l.tryAcquire(0)).toBe(false); // bucket empty
  });

  it('caps concurrency independently of the token bucket', () => {
    const l = new RpcRateLimiter({ burst: 100, perSec: 100, maxInFlight: 2 }, 0);
    expect(l.tryAcquire(0)).toBe(true); // in-flight 1
    expect(l.tryAcquire(0)).toBe(true); // in-flight 2
    expect(l.tryAcquire(0)).toBe(false); // concurrency cap hit despite tokens left
    l.release(); // in-flight 1
    expect(l.tryAcquire(0)).toBe(true); // slot freed
  });

  it('release never underflows below zero', () => {
    const l = new RpcRateLimiter({ burst: 1, perSec: 1, maxInFlight: 1 }, 0);
    l.release();
    l.release(); // no-op when nothing is in flight
    expect(l.tryAcquire(0)).toBe(true);
  });
});