Shade/packages/shade-files/tests/unit/rate-limiter.test.ts

import { describe, test, expect } from 'bun:test';
import {
  FsRateLimitError,
  QuotaExceededError,
  RateLimiter,
} from '../../src/index.js';

describe('RateLimiter — op bucket', () => {
  test('allows up to capacity then rejects', () => {
    const rl = new RateLimiter({ maxOpsPerMinutePerSender: 5, opCost: { default: 1 } });
    for (let i = 0; i < 5; i++) rl.acquire('alice', 'list');
    expect(() => rl.acquire('alice', 'list')).toThrow(FsRateLimitError);
  });

  test('rejection includes retryAfterMs', () => {
    const rl = new RateLimiter({ maxOpsPerMinutePerSender: 1, opCost: { default: 1 } });
    rl.acquire('alice', 'list');
    try {
      rl.acquire('alice', 'list');
      throw new Error('expected throw');
    } catch (err) {
      expect(err).toBeInstanceOf(FsRateLimitError);
      const payload = (err as FsRateLimitError).payload;
      expect(payload.retryAfterMs).toBeGreaterThan(0);
    }
  });

  test('different op costs respected', () => {
    const rl = new RateLimiter({
      maxOpsPerMinutePerSender: 10,
      opCost: { write: 5, default: 1 },
    });
    rl.acquire('alice', 'write'); // 10 - 5 = 5 left
    rl.acquire('alice', 'write'); // 5 - 5 = 0 left
    expect(() => rl.acquire('alice', 'write')).toThrow(FsRateLimitError);
  });

  test('per-sender isolation', () => {
    const rl = new RateLimiter({ maxOpsPerMinutePerSender: 1 });
    rl.acquire('alice', 'list');
    rl.acquire('bob', 'list'); // bob's bucket independent
    expect(() => rl.acquire('alice', 'list')).toThrow(FsRateLimitError);
    expect(() => rl.acquire('bob', 'list')).toThrow(FsRateLimitError);
  });

  test('release returns tokens', () => {
    const rl = new RateLimiter({ maxOpsPerMinutePerSender: 1 });
    rl.acquire('alice', 'list');
    rl.release('alice', 'list');
    rl.acquire('alice', 'list'); // should succeed again
  });
});

describe('RateLimiter — byte bucket', () => {
  test('quota exceeded triggers QuotaExceededError', () => {
    const rl = new RateLimiter({ maxBytesPerHourPerSender: 1000 });
    rl.acquire('alice', 'write', 600);
    rl.acquire('alice', 'write', 400);
    expect(() => rl.acquire('alice', 'write', 1)).toThrow(QuotaExceededError);
  });

  test('reconcile returns over-reserved bytes', () => {
    const rl = new RateLimiter({ maxBytesPerHourPerSender: 1000 });
    rl.acquire('alice', 'write', 800);
    rl.reconcile('alice', 800, 200); // we only used 200 of the 800 reserved
    rl.acquire('alice', 'write', 600); // capacity now 400 + 600 = 1000, fits 600
  });

  test('release returns reserved bytes', () => {
    const rl = new RateLimiter({ maxBytesPerHourPerSender: 100 });
    rl.acquire('alice', 'read', 80);
    rl.release('alice', 'read', 80);
    rl.acquire('alice', 'read', 80);
  });
});

describe('RateLimiter — refill', () => {
  test('tokens refill over time', async () => {
    const rl = new RateLimiter({
      // 6000/min = 100/sec = 0.1/ms; in 50ms we refill 5
      maxOpsPerMinutePerSender: 6000,
      opCost: { default: 1 },
    });
    // Drain by 5
    for (let i = 0; i < 5; i++) rl.acquire('alice', 'list');
    const before = rl.snapshot('alice')!.ops;
    await new Promise((r) => setTimeout(r, 50));
    const after = rl.snapshot('alice')!.ops;
    expect(after).toBeGreaterThan(before);
  });
});