/** * AAO Gateway Phase 2b — gateway_key_usage repository tests. * * Coverage: * - getGatewayKeyUsage returns null on miss, row on hit * - incrementGatewayKeyUsage UPSERTs first call, accumulates on second * - incrementGatewayKeyUsage clamps negative deltas to zero * - listGatewayKeyUsagesByKey orders by period_start DESC * - Cascade delete removes usage rows when key is deleted * - tokens_budget / rate_limit_rpm fields persist & round-trip * - updateGatewayVirtualKey patches fields independently and supports null reset */ import { describe, expect, it, beforeEach } from 'vitest'; import { Repository } from './repository.js'; import { hashKey } from '../gateway/key-format.js'; function makeRepo(): Repository { // ':memory:' triggers Repository.initSchema (creates all tables fresh). return new Repository(':memory:'); } function seedKey(repo: Repository, raw: string, team = 'alpha'): string { return repo.createGatewayVirtualKey({ keyHash: hashKey(raw), keyPrefix: raw.slice(0, 14), team, }).id; } describe('gateway_key_usage repository', () => { let repo: Repository; beforeEach(() => { repo = makeRepo(); }); it('getGatewayKeyUsage returns null when no row exists', () => { const id = seedKey(repo, 'sk-aao-test-1'); expect(repo.getGatewayKeyUsage(id, '2026-05')).toBeNull(); }); it('incrementGatewayKeyUsage creates row on first call, accumulates on second', () => { const id = seedKey(repo, 'sk-aao-test-2'); repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-05', tokensIn: 100, tokensOut: 50, requests: 1 }); let usage = repo.getGatewayKeyUsage(id, '2026-05'); expect(usage).not.toBeNull(); expect(usage!.tokensIn).toBe(100); expect(usage!.tokensOut).toBe(50); expect(usage!.requests).toBe(1); repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-05', tokensIn: 200, tokensOut: 80, requests: 1 }); usage = repo.getGatewayKeyUsage(id, '2026-05'); expect(usage!.tokensIn).toBe(300); expect(usage!.tokensOut).toBe(130); expect(usage!.requests).toBe(2); }); it('different periods get separate rows', () => { const id = seedKey(repo, 'sk-aao-test-3'); repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-04', tokensIn: 10, requests: 1 }); repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-05', tokensIn: 20, requests: 1 }); expect(repo.getGatewayKeyUsage(id, '2026-04')!.tokensIn).toBe(10); expect(repo.getGatewayKeyUsage(id, '2026-05')!.tokensIn).toBe(20); }); it('clamps negative deltas to zero', () => { const id = seedKey(repo, 'sk-aao-test-4'); repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-05', tokensIn: 100 }); repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-05', tokensIn: -500, tokensOut: -10, requests: -1 }); const usage = repo.getGatewayKeyUsage(id, '2026-05')!; expect(usage.tokensIn).toBe(100); // unchanged expect(usage.tokensOut).toBe(0); expect(usage.requests).toBe(0); }); it('listGatewayKeyUsagesByKey returns rows newest period first', () => { const id = seedKey(repo, 'sk-aao-test-5'); repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-03', requests: 1 }); repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-05', requests: 1 }); repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-04', requests: 1 }); const list = repo.listGatewayKeyUsagesByKey(id); expect(list.map(u => u.periodStart)).toEqual(['2026-05', '2026-04', '2026-03']); }); it('listGatewayKeyUsagesByKey honors limit', () => { const id = seedKey(repo, 'sk-aao-test-6'); for (const m of ['2026-01', '2026-02', '2026-03', '2026-04']) { repo.incrementGatewayKeyUsage({ keyId: id, period: m, requests: 1 }); } const list = repo.listGatewayKeyUsagesByKey(id, { limit: 2 }); expect(list).toHaveLength(2); expect(list[0]!.periodStart).toBe('2026-04'); }); it('cascade-deletes usage rows when key is hard-deleted', () => { const id = seedKey(repo, 'sk-aao-test-7'); repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-05', tokensIn: 1 }); expect(repo.getGatewayKeyUsage(id, '2026-05')).not.toBeNull(); // The Repository's defense-in-depth guard refuses to delete // config-import rows. Seeded above as 'admin' so this is legal. repo.deleteGatewayVirtualKey(id); expect(repo.getGatewayKeyUsage(id, '2026-05')).toBeNull(); }); }); describe('gateway_virtual_keys budget/rate fields', () => { let repo: Repository; beforeEach(() => { repo = makeRepo(); }); it('round-trips tokensBudget and rateLimitRpm on create', () => { const created = repo.createGatewayVirtualKey({ keyHash: hashKey('sk-aao-create-budget'), keyPrefix: 'sk-aao-create', team: 'alpha', tokensBudget: 1_000_000, rateLimitRpm: 60, }); expect(created.tokensBudget).toBe(1_000_000); expect(created.rateLimitRpm).toBe(60); const refreshed = repo.findGatewayVirtualKeyById(created.id)!; expect(refreshed.tokensBudget).toBe(1_000_000); expect(refreshed.rateLimitRpm).toBe(60); }); it('defaults to null when budget/rate omitted', () => { const created = repo.createGatewayVirtualKey({ keyHash: hashKey('sk-aao-no-limits'), keyPrefix: 'sk-aao-no-lim', team: 'beta', }); expect(created.tokensBudget).toBeNull(); expect(created.rateLimitRpm).toBeNull(); }); it('coerces zero or negative limits to null (defensive)', () => { const created = repo.createGatewayVirtualKey({ keyHash: hashKey('sk-aao-bad-limits'), keyPrefix: 'sk-aao-bad', team: 'gamma', tokensBudget: 0, rateLimitRpm: -5, }); expect(created.tokensBudget).toBeNull(); expect(created.rateLimitRpm).toBeNull(); }); it('updateGatewayVirtualKey patches only specified fields', () => { const created = repo.createGatewayVirtualKey({ keyHash: hashKey('sk-aao-patch'), keyPrefix: 'sk-aao-patch', team: 'alpha', tokensBudget: 1000, rateLimitRpm: 60, allowedModels: ['qwen3:8b'], }); // Patch only tokensBudget — other fields untouched const after1 = repo.updateGatewayVirtualKey(created.id, { tokensBudget: 5000 }); expect(after1.tokensBudget).toBe(5000); expect(after1.rateLimitRpm).toBe(60); expect(after1.allowedModels).toEqual(['qwen3:8b']); // Patch allowedModels alone const after2 = repo.updateGatewayVirtualKey(created.id, { allowedModels: ['qwen3:14b'] }); expect(after2.allowedModels).toEqual(['qwen3:14b']); expect(after2.tokensBudget).toBe(5000); // Reset rate limit to null (unlimited) const after3 = repo.updateGatewayVirtualKey(created.id, { rateLimitRpm: null }); expect(after3.rateLimitRpm).toBeNull(); // Reset allowedModels to null (no allowlist) const after4 = repo.updateGatewayVirtualKey(created.id, { allowedModels: null }); expect(after4.allowedModels).toBeNull(); }); it('updateGatewayVirtualKey throws for unknown id', () => { expect(() => repo.updateGatewayVirtualKey('does-not-exist', { tokensBudget: 1 })).toThrow(); }); });