85 lines
3.9 KiB
TypeScript
85 lines
3.9 KiB
TypeScript
import {
|
|
EMBEDDING_LIMITS,
|
|
EMBEDDING_QUOTA,
|
|
embeddingUnits,
|
|
firstOversizedChunk,
|
|
planEmbeddingBatches,
|
|
refundEmbeddingUnits,
|
|
reserveEmbeddingUnits,
|
|
usageWindowStart,
|
|
} from './embedding-quota.ts'
|
|
import { MemoryUsageStore } from './embedding-quota.fake.ts'
|
|
|
|
function assert(condition: boolean, message: string): asserts condition {
|
|
if (!condition) throw new Error(message)
|
|
}
|
|
|
|
const NOW = new Date('2026-09-28T12:00:00Z')
|
|
|
|
Deno.test('units count every started 1,000 characters, at least one per input', () => {
|
|
assert(embeddingUnits(['a']) === 1, 'short input')
|
|
assert(embeddingUnits(['x'.repeat(1_000), 'x'.repeat(1_001)]) === 3, 'boundary rounding')
|
|
})
|
|
|
|
Deno.test('oversized chunks are found before any spend', () => {
|
|
const ok = { content: 'x'.repeat(EMBEDDING_LIMITS.maxChunkChars) }
|
|
const big = { content: 'x'.repeat(EMBEDDING_LIMITS.maxChunkChars + 1) }
|
|
assert(firstOversizedChunk([ok, ok]) === -1, 'at-limit chunk rejected')
|
|
assert(firstOversizedChunk([ok, big]) === 1, 'oversized chunk missed')
|
|
})
|
|
|
|
Deno.test('batches stay under the input and character caps', () => {
|
|
const small = Array.from({ length: 250 }, (_, i) => ({ content: `c${i}` }))
|
|
assert(planEmbeddingBatches(small).map((b) => b.length).join(',') === '100,100,50', 'input cap')
|
|
const large = Array.from({ length: 30 }, () => ({ content: 'x'.repeat(8_000) }))
|
|
const batches = planEmbeddingBatches(large)
|
|
assert(batches.every((b) => b.reduce((s, c) => s + c.content.length, 0) <= EMBEDDING_LIMITS.maxBatchChars), 'char cap')
|
|
assert(batches.flat().length === 30, 'chunks lost')
|
|
})
|
|
|
|
Deno.test('weekly window covers the six days before today', () => {
|
|
assert(usageWindowStart('weekly', NOW) === '2026-09-22', 'weekly start')
|
|
assert(usageWindowStart('daily', NOW) === null, 'daily has no earlier window')
|
|
})
|
|
|
|
Deno.test('free tier is capped weekly and a denied reservation is rolled back', async () => {
|
|
const store = new MemoryUsageStore('free', '2026-09-28')
|
|
store.rows.set('2026-09-25', EMBEDDING_QUOTA.free.limit - 10)
|
|
const denied = await reserveEmbeddingUnits(store, 'u1', 11, NOW)
|
|
assert(!denied.allowed, 'over-limit reservation was allowed')
|
|
assert((store.rows.get('2026-09-28') ?? 0) === 0, 'denied units were kept')
|
|
const allowed = await reserveEmbeddingUnits(store, 'u1', 10, NOW)
|
|
assert(allowed.allowed && allowed.current === EMBEDDING_QUOTA.free.limit, 'exact fit denied')
|
|
})
|
|
|
|
Deno.test('usage older than the window does not count', async () => {
|
|
const store = new MemoryUsageStore('free', '2026-09-28')
|
|
store.rows.set('2026-09-21', EMBEDDING_QUOTA.free.limit)
|
|
assert((await reserveEmbeddingUnits(store, 'u1', 5, NOW)).allowed, 'stale usage counted')
|
|
})
|
|
|
|
Deno.test('enterprise is unlimited but still counted', async () => {
|
|
const store = new MemoryUsageStore('enterprise', '2026-09-28')
|
|
const result = await reserveEmbeddingUnits(store, 'u1', 1_000_000, NOW)
|
|
assert(result.allowed && store.rows.get('2026-09-28') === 1_000_000, 'unlimited tier not counted')
|
|
})
|
|
|
|
Deno.test('refunds are capped by the reservation and skipped after the day rolls over', async () => {
|
|
const store = new MemoryUsageStore('pro', '2026-09-28')
|
|
const reservation = await reserveEmbeddingUnits(store, 'u1', 40, NOW)
|
|
await refundEmbeddingUnits(store, 'u1', reservation, 1_000, NOW)
|
|
assert(store.rows.get('2026-09-28') === 0, 'refund exceeded the reserved units')
|
|
|
|
const again = await reserveEmbeddingUnits(store, 'u1', 40, NOW)
|
|
store.today = '2026-09-29'
|
|
await refundEmbeddingUnits(store, 'u1', again, 40, new Date('2026-09-29T00:00:01Z'))
|
|
assert(!store.rows.has('2026-09-29'), 'next-day refund credited the new day')
|
|
})
|
|
|
|
Deno.test('a store failure during refund is swallowed', async () => {
|
|
const store = new MemoryUsageStore('pro', '2026-09-28')
|
|
const reservation = await reserveEmbeddingUnits(store, 'u1', 3, NOW)
|
|
store.failIncrement = true
|
|
await refundEmbeddingUnits(store, 'u1', reservation, 3, NOW)
|
|
assert(store.rows.get('2026-09-28') === 3, 'usage changed')
|
|
})
|