Spaces:
Running
Running
File size: 2,223 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | import { vi } from "vitest";
/**
* Common mock setup for text generation tests
*/
const createCommonMocks = () => ({
pubSub: {
getSettings: vi.fn(),
updateModelLoadingProgress: vi.fn(),
updateModelSizeInMegabytes: vi.fn(),
updateTextGenerationState: vi.fn(),
updateResponse: vi.fn(),
canStartResponding: vi.fn(),
},
textGenerationUtilities: {
canStartResponding: vi.fn(),
defaultContextSize: 128,
getDefaultChatMessages: vi.fn(() => []),
getFormattedSearchResults: vi.fn(() => []),
getDefaultChatCompletionCreateParamsStreaming: vi.fn(),
},
});
/**
* Setup common vi.mocks for text generation tests
*/
export const setupCommonTextGenerationMocks = () => {
vi.mock("@root/package.json", () => ({
repository: { url: "https://github.com/owner/repo" },
version: "1.0.0",
}));
};
/**
* Mock pubSub module with specific settings
*/
export const mockPubSub = (settings: Record<string, unknown> = {}) => {
const commonMocks = createCommonMocks();
return {
...commonMocks.pubSub,
getSettings: vi.fn(() => settings),
getQuery: vi.fn(() => "test query"),
getQuerySuggestions: vi.fn(() => []),
updateQuerySuggestions: vi.fn(),
queryPubSub: [vi.fn(), vi.fn(), vi.fn()],
settingsPubSub: [vi.fn(), vi.fn(), vi.fn()],
updateTextGenerationState: vi.fn(),
getTextGenerationState: vi.fn(() => "idle"),
};
};
/**
* Mock textGenerationUtilities module
*/
export const mockTextGenerationUtilities = (
overrides: Record<string, unknown> = {},
) => {
const commonMocks = createCommonMocks();
return {
...commonMocks.textGenerationUtilities,
...overrides,
};
};
/**
* Helper function to test common text generation behavior
*/
export const testTextGenerationBehavior = async (
generateFunction: () => Promise<void>,
expectedResponse: string,
) => {
await generateFunction();
const utils = await import("./textGenerationUtilities");
const pubSub = await import("./pubSub");
expect(utils.canStartResponding).toHaveBeenCalled();
expect(pubSub.updateTextGenerationState).toHaveBeenCalledWith(
"preparingToGenerate",
);
expect(pubSub.updateResponse).toHaveBeenCalledWith(expectedResponse);
};
|