File size: 2,589 Bytes
2b64d42 | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | import { afterEach, describe, it } from 'node:test';
import assert from 'node:assert/strict';
import { addAccountByKey, removeAccount } from '../src/auth.js';
import { cacheClear, cacheKey, cacheSet } from '../src/cache.js';
import { handleChatCompletions } from '../src/handlers/chat.js';
const createdAccountIds = [];
function fakeRes() {
const listeners = new Map();
return {
body: '',
writableEnded: false,
write(chunk) {
this.body += String(chunk);
return true;
},
end(chunk) {
if (chunk) this.write(chunk);
this.writableEnded = true;
for (const cb of listeners.get('close') || []) cb();
},
on(event, cb) {
if (!listeners.has(event)) listeners.set(event, []);
listeners.get(event).push(cb);
return this;
},
};
}
function parseChatFrames(raw) {
return raw
.split('\n\n')
.filter(Boolean)
.filter(frame => !frame.startsWith(':'))
.map(frame => {
const dataLine = frame.split('\n').find(line => line.startsWith('data: '));
const payload = dataLine?.slice(6) || '';
return payload === '[DONE]' ? '[DONE]' : JSON.parse(payload);
});
}
afterEach(() => {
cacheClear();
while (createdAccountIds.length) {
removeAccount(createdAccountIds.pop());
}
});
describe('chat cache-hit stream shape', () => {
it('matches the live-stream finish chunk plus terminal usage chunk shape', async () => {
const account = addAccountByKey(`cache-key-${Date.now()}`, 'cache-hit');
createdAccountIds.push(account.id);
const body = {
model: 'gemini-2.5-flash',
stream: true,
messages: [{ role: 'user', content: 'hi' }],
};
cacheSet(cacheKey(body), { text: 'cached answer', thinking: 'cached thinking' });
const result = await handleChatCompletions(body);
assert.equal(result.status, 200);
assert.equal(result.stream, true);
const res = fakeRes();
await result.handler(res);
const frames = parseChatFrames(res.body);
const finishChunk = frames.at(-3);
const usageChunk = frames.at(-2);
assert.equal(finishChunk.choices[0].finish_reason, 'stop');
assert.equal('usage' in finishChunk, false);
assert.deepEqual(usageChunk.choices, []);
assert.deepEqual(usageChunk.usage, {
cached: true,
prompt_tokens: 1,
completion_tokens: 4,
total_tokens: 5,
input_tokens: 1,
output_tokens: 4,
prompt_tokens_details: { cached_tokens: 1 },
completion_tokens_details: { reasoning_tokens: 0 },
});
assert.equal(frames.at(-1), '[DONE]');
});
});
|