rate-limit.test.ts

rate-limit.test.tsapps/realtime/test/rate-limit.test.tsraw ← test3,700 B← Back to the app
1import { describe, it, expect, beforeEach } from "vitest";
2import { TokenBucket, ConnectLimiter } from "../src/rate-limit.js";
3 
4/** A deterministic fake clock. */
5function fakeClock(initial = 0) {
6  let t = initial;
7  return {
8    now: () => t,
9    advance: (ms: number) => {
10      t += ms;
11    },
12  };
13}
14 
15describe("TokenBucket (per-socket sustained rate, R4)", () => {
16  let clock: ReturnType<typeof fakeClock>;
17  let bucket: TokenBucket;
18 
19  beforeEach(() => {
20    clock = fakeClock();
21    bucket = new TokenBucket({ capacity: 5, rate: 5, now: clock.now });
22  });
23 
24  it("serves the full burst capacity without consuming below zero", () => {
25    for (let i = 0; i < 5; i++) {
26      expect(bucket.consume("a").ok).toBe(true);
27    }
28  });
29 
30  it("rejects the request past capacity with a positive retryAfterMs", () => {
31    for (let i = 0; i < 5; i++) bucket.consume("a");
32    const denied = bucket.consume("a");
33    expect(denied.ok).toBe(false);
34    if (!denied.ok) expect(denied.retryAfterMs).toBeGreaterThan(0);
35  });
36 
37  it("tracks keys independently", () => {
38    for (let i = 0; i < 5; i++) bucket.consume("a");
39    expect(bucket.consume("a").ok).toBe(false);
40    expect(bucket.consume("b").ok).toBe(true);
41  });
42 
43  it("refills continuously — the bucket is usable again after (capacity/rate) seconds", () => {
44    for (let i = 0; i < 5; i++) bucket.consume("a");
45    expect(bucket.consume("a").ok).toBe(false);
46    // 5 capacity / 5 per s → a full refill after 1000 ms.
47    clock.advance(1000);
48    expect(bucket.consume("a").ok).toBe(true);
49  });
50 
51  it("refills are capped at the capacity (no unbounded accumulation)", () => {
52    clock.advance(60_000); // an hour of idleness
53    expect(bucket.available("a")).toBe(5); // exactly the capacity, not 3000
54  });
55 
56  it("drop removes a bucket and size reflects live buckets", () => {
57    bucket.consume("a");
58    bucket.consume("b");
59    expect(bucket.size()).toBe(2);
60    bucket.drop("a");
61    expect(bucket.size()).toBe(1);
62  });
63 
64  it("sweep drops buckets idle longer than staleMs, keeps active ones", () => {
65    bucket.consume("idle");
66    clock.advance(10_000);
67    bucket.consume("fresh");
68    expect(bucket.sweep(5000)).toBe(1); // only "idle" is stale
69    expect(bucket.available("fresh")).toBeGreaterThan(0);
70  });
71});
72 
73describe("ConnectLimiter (connect-time admission, R4)", () => {
74  let clock: ReturnType<typeof fakeClock>;
75  let limiter: ConnectLimiter;
76 
77  beforeEach(() => {
78    clock = fakeClock();
79    limiter = new ConnectLimiter({ windowMs: 60_000, max: 3, now: clock.now });
80  });
81 
82  it("admits up to the max per window, then rejects", () => {
83    expect(limiter.tryAdmit("1.2.3.4").ok).toBe(true);
84    expect(limiter.tryAdmit("1.2.3.4").ok).toBe(true);
85    expect(limiter.tryAdmit("1.2.3.4").ok).toBe(true);
86    expect(limiter.tryAdmit("1.2.3.4")).toEqual({ ok: false });
87  });
88 
89  it("tracks keys (IPs) independently", () => {
90    limiter.tryAdmit("a");
91    limiter.tryAdmit("a");
92    limiter.tryAdmit("a");
93    expect(limiter.tryAdmit("a").ok).toBe(false);
94    expect(limiter.tryAdmit("b").ok).toBe(true);
95  });
96 
97  it("admits again once the sliding window has passed", () => {
98    limiter.tryAdmit("a");
99    limiter.tryAdmit("a");
100    limiter.tryAdmit("a");
101    expect(limiter.tryAdmit("a").ok).toBe(false);
102    clock.advance(60_001);
103    expect(limiter.tryAdmit("a").ok).toBe(true);
104  });
105 
106  it("count reflects the in-window hits", () => {
107    limiter.tryAdmit("a");
108    limiter.tryAdmit("a");
109    expect(limiter.count("a")).toBe(2);
110  });
111 
112  it("sweep drops idle keys", () => {
113    limiter.tryAdmit("a");
114    clock.advance(120_000);
115    expect(limiter.sweep(60_000)).toBe(1);
116    expect(limiter.size()).toBe(0);
117  });
118});
119