From 68a0ef664819c9d2755ba3eaa369b3be182a3ce8 Mon Sep 17 00:00:00 2001
From: Alex Gleason <alex@alexgleason.me>
Date: Sat, 25 Jan 2025 15:20:52 -0600
Subject: [PATCH] Add ratelimiter tests

---
 .../ratelimiter/MemoryRateLimiter.test.ts     | 31 +++++++++++++
 src/utils/ratelimiter/MemoryRateLimiter.ts    |  2 +-
 .../ratelimiter/MultiRateLimiter.test.ts      | 39 ++++++++++++++++
 src/utils/ratelimiter/MultiRateLimiter.ts     | 45 +++++++++++++++++++
 4 files changed, 116 insertions(+), 1 deletion(-)
 create mode 100644 src/utils/ratelimiter/MemoryRateLimiter.test.ts
 create mode 100644 src/utils/ratelimiter/MultiRateLimiter.test.ts
 create mode 100644 src/utils/ratelimiter/MultiRateLimiter.ts

diff --git a/src/utils/ratelimiter/MemoryRateLimiter.test.ts b/src/utils/ratelimiter/MemoryRateLimiter.test.ts
new file mode 100644
index 00000000..2da6b2d1
--- /dev/null
+++ b/src/utils/ratelimiter/MemoryRateLimiter.test.ts
@@ -0,0 +1,31 @@
+import { assertEquals, assertThrows } from '@std/assert';
+
+import { MemoryRateLimiter } from './MemoryRateLimiter.ts';
+import { RateLimitError } from './RateLimitError.ts';
+
+Deno.test('MemoryRateLimiter', async (t) => {
+  const limit = 5;
+  const window = 100;
+
+  using limiter = new MemoryRateLimiter({ limit, window });
+
+  await t.step('can hit up to limit', () => {
+    for (let i = 0; i < limit; i++) {
+      const client = limiter.client('test');
+      assertEquals(client.hits, i);
+      client.hit();
+    }
+  });
+
+  await t.step('throws when hit if limit exceeded', () => {
+    assertThrows(() => limiter.client('test').hit(), RateLimitError);
+  });
+
+  await t.step('can hit after window resets', async () => {
+    await new Promise((resolve) => setTimeout(resolve, window + 1));
+
+    const client = limiter.client('test');
+    assertEquals(client.hits, 0);
+    client.hit();
+  });
+});
diff --git a/src/utils/ratelimiter/MemoryRateLimiter.ts b/src/utils/ratelimiter/MemoryRateLimiter.ts
index b3f14d81..0eaa5540 100644
--- a/src/utils/ratelimiter/MemoryRateLimiter.ts
+++ b/src/utils/ratelimiter/MemoryRateLimiter.ts
@@ -35,7 +35,7 @@ export class MemoryRateLimiter implements RateLimiter {
       return curr;
     }
 
-    if (prev) {
+    if (prev && prev.resetAt > new Date()) {
       this.current.set(key, prev);
       this.previous.delete(key);
       return prev;
diff --git a/src/utils/ratelimiter/MultiRateLimiter.test.ts b/src/utils/ratelimiter/MultiRateLimiter.test.ts
new file mode 100644
index 00000000..3cfa4696
--- /dev/null
+++ b/src/utils/ratelimiter/MultiRateLimiter.test.ts
@@ -0,0 +1,39 @@
+import { assertEquals, assertThrows } from '@std/assert';
+
+import { MemoryRateLimiter } from './MemoryRateLimiter.ts';
+import { MultiRateLimiter } from './MultiRateLimiter.ts';
+
+Deno.test('MultiRateLimiter', async (t) => {
+  using limiter1 = new MemoryRateLimiter({ limit: 5, window: 100 });
+  using limiter2 = new MemoryRateLimiter({ limit: 8, window: 200 });
+
+  const limiter = new MultiRateLimiter([limiter1, limiter2]);
+
+  await t.step('can hit up to first limit', () => {
+    for (let i = 0; i < limiter1.limit; i++) {
+      const client = limiter.client('test');
+      assertEquals(client.hits, i);
+      client.hit();
+    }
+  });
+
+  await t.step('throws when hit if first limit exceeded', () => {
+    assertThrows(() => limiter.client('test').hit(), Error);
+  });
+
+  await t.step('can hit up to second limit after the first window resets', async () => {
+    await new Promise((resolve) => setTimeout(resolve, limiter1.window + 1));
+
+    const limit = limiter2.limit - limiter1.limit - 1;
+
+    for (let i = 0; i < limit; i++) {
+      const client = limiter.client('test');
+      assertEquals(client.hits, i);
+      client.hit();
+    }
+  });
+
+  await t.step('throws when hit if second limit exceeded', () => {
+    assertThrows(() => limiter.client('test').hit(), Error);
+  });
+});
diff --git a/src/utils/ratelimiter/MultiRateLimiter.ts b/src/utils/ratelimiter/MultiRateLimiter.ts
new file mode 100644
index 00000000..dc9b62a7
--- /dev/null
+++ b/src/utils/ratelimiter/MultiRateLimiter.ts
@@ -0,0 +1,45 @@
+import { RateLimiter, RateLimiterClient } from './types.ts';
+
+export class MultiRateLimiter {
+  constructor(private limiters: RateLimiter[]) {}
+
+  client(key: string): RateLimiterClient {
+    return new MultiRateLimiterClient(key, this.limiters);
+  }
+}
+
+class MultiRateLimiterClient implements RateLimiterClient {
+  constructor(private key: string, private limiters: RateLimiter[]) {
+    if (!limiters.length) {
+      throw new Error('No limiters provided');
+    }
+  }
+
+  get hits(): number {
+    return this.limiters[0].client(this.key).hits;
+  }
+
+  get resetAt(): Date {
+    return this.limiters[0].client(this.key).resetAt;
+  }
+
+  get remaining(): number {
+    return this.limiters[0].client(this.key).remaining;
+  }
+
+  hit(n?: number): void {
+    let error: unknown;
+
+    for (const limiter of this.limiters) {
+      try {
+        limiter.client(this.key).hit(n);
+      } catch (e) {
+        error ??= e;
+      }
+    }
+
+    if (error instanceof Error) {
+      throw error;
+    }
+  }
+}