122 lines
4.6 KiB
TypeScript
122 lines
4.6 KiB
TypeScript
import { describe, expect, test } from "bun:test";
|
|
import { fetchModelFile, type ModelFileCacheStorage } from "../src/worker/modelFiles";
|
|
import type { LoadProgressMessage } from "../src/worker/protocol";
|
|
|
|
// Cache-API-backed weight download for the Parakeet path (D-WS7a, D-005).
|
|
|
|
const URL_A = "https://huggingface.co/org/model/resolve/rev/encoder.onnx";
|
|
|
|
function fakeCaches(initial: Record<string, Uint8Array> = {}) {
|
|
const stored = new Map<string, Uint8Array>(Object.entries(initial));
|
|
const opened: string[] = [];
|
|
const storage: ModelFileCacheStorage = {
|
|
async open(name) {
|
|
opened.push(name);
|
|
return {
|
|
async match(url) {
|
|
const bytes = stored.get(url);
|
|
return bytes ? new Response(bytes as Uint8Array<ArrayBuffer>) : undefined;
|
|
},
|
|
async put(url, response) {
|
|
stored.set(url, new Uint8Array(await response.arrayBuffer()));
|
|
},
|
|
};
|
|
},
|
|
};
|
|
return { storage, stored, opened };
|
|
}
|
|
|
|
function streamingFetch(bytes: Uint8Array, pieces: number, status = 200) {
|
|
const calls: string[] = [];
|
|
const fetchFn = async (url: string) => {
|
|
calls.push(url);
|
|
if (status !== 200) return new Response("nope", { status });
|
|
const size = Math.ceil(bytes.length / pieces);
|
|
const body = new ReadableStream<Uint8Array>({
|
|
start(controller) {
|
|
for (let i = 0; i < bytes.length; i += size) controller.enqueue(bytes.slice(i, i + size));
|
|
controller.close();
|
|
},
|
|
});
|
|
return new Response(body, { headers: { "Content-Length": String(bytes.length) } });
|
|
};
|
|
return { fetchFn, calls };
|
|
}
|
|
|
|
describe("fetchModelFile", () => {
|
|
test("a cached file is returned without touching the network", async () => {
|
|
const cached = new Uint8Array([1, 2, 3]);
|
|
const { storage, opened } = fakeCaches({ [URL_A]: cached });
|
|
const { fetchFn, calls } = streamingFetch(new Uint8Array(), 1);
|
|
const events: LoadProgressMessage[] = [];
|
|
const bytes = await fetchModelFile(URL_A, "encoder.onnx", {
|
|
caches: storage,
|
|
fetch: fetchFn,
|
|
onProgress: (event) => events.push(event),
|
|
});
|
|
expect([...bytes]).toEqual([1, 2, 3]);
|
|
expect(calls).toEqual([]);
|
|
expect(opened).toEqual(["transformers-cache"]);
|
|
expect(events.map((event) => event.status)).toEqual(["initiate", "done"]);
|
|
});
|
|
|
|
test("a miss downloads with byte progress and stores the complete file", async () => {
|
|
const payload = new Uint8Array(1000).map((_, i) => i % 251);
|
|
const { storage, stored } = fakeCaches();
|
|
const { fetchFn, calls } = streamingFetch(payload, 4);
|
|
const events: LoadProgressMessage[] = [];
|
|
const bytes = await fetchModelFile(URL_A, "encoder.onnx", {
|
|
caches: storage,
|
|
fetch: fetchFn,
|
|
onProgress: (event) => events.push(event),
|
|
});
|
|
expect(calls).toEqual([URL_A]);
|
|
expect([...bytes]).toEqual([...payload]);
|
|
expect([...stored.get(URL_A)!]).toEqual([...payload]);
|
|
const progress = events.filter((event) => event.status === "progress");
|
|
expect(progress.length).toBe(4);
|
|
const last = progress[progress.length - 1]!;
|
|
expect(last).toMatchObject({ file: "encoder.onnx", loaded: 1000, total: 1000, progress: 100 });
|
|
expect(events[0]!.status).toBe("initiate");
|
|
expect(events[events.length - 1]!.status).toBe("done");
|
|
});
|
|
|
|
test("an HTTP error throws a plain message and stores nothing", async () => {
|
|
const { storage, stored } = fakeCaches();
|
|
const { fetchFn } = streamingFetch(new Uint8Array(), 1, 404);
|
|
await expect(
|
|
fetchModelFile(URL_A, "encoder.onnx", { caches: storage, fetch: fetchFn, onProgress: () => {} }),
|
|
).rejects.toThrow("Could not download encoder.onnx (HTTP 404)");
|
|
expect(stored.size).toBe(0);
|
|
});
|
|
|
|
test("an interrupted download stores nothing (no false 'available offline')", async () => {
|
|
const { storage, stored } = fakeCaches();
|
|
const fetchFn = async () =>
|
|
new Response(
|
|
new ReadableStream<Uint8Array>({
|
|
start(controller) {
|
|
controller.enqueue(new Uint8Array(10));
|
|
controller.error(new Error("network down"));
|
|
},
|
|
}),
|
|
{ headers: { "Content-Length": "100" } },
|
|
);
|
|
await expect(
|
|
fetchModelFile(URL_A, "encoder.onnx", { caches: storage, fetch: fetchFn, onProgress: () => {} }),
|
|
).rejects.toThrow();
|
|
expect(stored.size).toBe(0);
|
|
});
|
|
|
|
test("without Cache Storage it still downloads", async () => {
|
|
const { fetchFn, calls } = streamingFetch(new Uint8Array([9]), 1);
|
|
const bytes = await fetchModelFile(URL_A, "encoder.onnx", {
|
|
caches: undefined,
|
|
fetch: fetchFn,
|
|
onProgress: () => {},
|
|
});
|
|
expect([...bytes]).toEqual([9]);
|
|
expect(calls).toEqual([URL_A]);
|
|
});
|
|
});
|