Spaces:
Running
Running
File size: 1,672 Bytes
10d1fd4 | 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 | import { beforeEach, describe, it, vi } from "vitest";
import {
mockPubSub,
mockTextGenerationUtilities,
testTextGenerationBehavior,
} from "./testUtils";
vi.stubGlobal("Worker", vi.fn());
vi.mock("./pubSub", () =>
mockPubSub({
enableAiResponse: true,
webLlmModelId: "model-id",
}),
);
vi.mock("./textGenerationUtilities", () =>
mockTextGenerationUtilities({
handleStreamingResponse: vi.fn(async (_completion, onUpdate) => {
await onUpdate("partial");
return "partial";
}),
}),
);
vi.mock("@mlc-ai/web-llm", () => ({
CreateWebWorkerMLCEngine: vi.fn().mockResolvedValue({
chat: {
completions: {
create: vi.fn().mockResolvedValue({ fullStream: [] }),
},
},
resetChat: vi.fn(),
runtimeStatsText: vi.fn().mockResolvedValue("stats"),
unload: vi.fn(),
}),
CreateMLCEngine: vi.fn(),
hasModelInCache: vi.fn().mockResolvedValue(true),
prebuiltAppConfig: {
model_list: [{ model_id: "model-id", vram_required_MB: 1024 }],
},
}));
import * as mod from "./textGenerationWithWebLlm";
describe("generateTextWithWebLlm", () => {
beforeEach(() => {
vi.resetAllMocks();
// biome-ignore lint/suspicious/noExplicitAny: Necessary for test mocking
(mod as any).initializeWebLlmEngine = vi.fn().mockResolvedValue({
chat: {
completions: { create: vi.fn().mockResolvedValue("completion") },
},
runtimeStatsText: vi.fn().mockResolvedValue("stats"),
unload: vi.fn(),
});
});
it("calls helpers and updates state", async () => {
await testTextGenerationBehavior(
() => mod.generateTextWithWebLlm(),
"partial",
);
});
});
|