memory-store.test.js
1.96 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
import { describe, it, expect } from 'vitest'
import { createMemoryStore } from '../src/memory-store.js'
const rule = { endpoint: 'llm', by: 'ip', capacity: 2, refillPerSec: 1 }
describe('memory-store', () => {
it('take 在 capacity 内连续放行、超限拒绝', async () => {
const store = createMemoryStore()
const a = await store.take('k', 1, rule, 0)
const b = await store.take('k', 1, rule, 0)
const c = await store.take('k', 1, rule, 0)
expect(a.allowed).toBe(true)
expect(b.allowed).toBe(true)
expect(c.allowed).toBe(false)
expect(c.retryAfterMs).toBeGreaterThan(0)
})
it('时间推进后恢复放行', async () => {
const store = createMemoryStore()
await store.take('k', 1, rule, 0)
await store.take('k', 1, rule, 0)
expect((await store.take('k', 1, rule, 0)).allowed).toBe(false)
// 1s 后补 1 个
expect((await store.take('k', 1, rule, 1000)).allowed).toBe(true)
})
it('peek 不消耗令牌', async () => {
const store = createMemoryStore()
await store.take('k', 1, rule, 0) // 剩 1
const p1 = await store.peek('k', rule, 0)
const p2 = await store.peek('k', rule, 0)
expect(p1.remaining).toBe(1)
expect(p2.remaining).toBe(1)
// peek 后仍能再 take 1 次成功
expect((await store.take('k', 1, rule, 0)).allowed).toBe(true)
})
it('并发顺序调用不超发（同步原子，不超过 capacity）', async () => {
const store = createMemoryStore()
const results = await Promise.all(
Array.from({ length: 5 }, () => store.take('k', 1, rule, 0)),
)
const allowed = results.filter((r) => r.allowed).length
expect(allowed).toBe(2) // 恰好 capacity 个
})
it('不同 key 互不串桶', async () => {
const store = createMemoryStore()
await store.take('a', 1, rule, 0)
await store.take('a', 1, rule, 0)
expect((await store.take('a', 1, rule, 0)).allowed).toBe(false)
expect((await store.take('b', 1, rule, 0)).allowed).toBe(true)
})
})