23 lines
1 KiB
TypeScript
23 lines
1 KiB
TypeScript
import { describe, expect, it } from "vitest";
|
|
import { InferenceQueue } from "../src/inference/queue";
|
|
import type { InferenceRequest } from "../src/shared/types";
|
|
|
|
const request = (id: string, priority: 0 | 1 | 2 | 3): InferenceRequest => ({
|
|
id, text: `text ${id}`, site: "reddit", priority, requestId: id, navigationId: "nav"
|
|
});
|
|
|
|
describe("InferenceQueue", () => {
|
|
it("batches requests and caches equivalent text", async () => {
|
|
let calls = 0;
|
|
const queue = new InferenceQueue(async (requests) => {
|
|
calls += 1;
|
|
return requests.map((item) => ({ requestId: item.requestId, id: item.id, textHash: "hash", label: "toxic" as const, probability: .9, navigationId: item.navigationId }));
|
|
}, undefined, 10, 8, 0);
|
|
const first = await queue.enqueue(request("a", 0));
|
|
const second = await queue.enqueue({ ...request("b", 0), text: "text a" });
|
|
expect(first.probability).toBe(.9);
|
|
expect(second.requestId).toBe("b");
|
|
expect(calls).toBe(1);
|
|
expect(queue.cacheSize).toBe(1);
|
|
});
|
|
});
|