@competent-devs/test-forge/dist/src/draft-test-forge/test-forge.js

Package for UI unit test generation based on storybook context

import * as path from "node:path";
import * as fs from "node:fs/promises";
import { globbyStream } from "globby";
import Mustache from "mustache";
import * as prettier from "prettier";
import { extractSource, readCsf } from "storybook/internal/csf-tools";
import { loadMainConfig } from "storybook/internal/common";
import { HumanMessage } from "@langchain/core/messages";
import { invokeLlm } from "./agent.js";
const ROOT_DIR = path.resolve("./");
const STORYBOOK_CONFIG_DIR = path.resolve(ROOT_DIR, "./.storybook");
const EXTENSION_MDX = ".mdx";
export const getPromptTemplate = async () => {
    return await fs.readFile(path.resolve("./src/test-forge/prompt.mustache"), "utf8");
};
export const getStorybookConfig = async () => {
    return await loadMainConfig({
        configDir: STORYBOOK_CONFIG_DIR,
    });
};
export const testForge = async () => {
    const storybookConfig = await getStorybookConfig();
    for await (const filePath of globbyStream(storybookConfig.stories, {
        onlyFiles: true,
        cwd: STORYBOOK_CONFIG_DIR,
    })) {
        if (isMDX(filePath)) {
            continue;
        }
        console.group(`===== ${filePath} =====`);
        const relativepathtoroot = path.relative(STORYBOOK_CONFIG_DIR, ROOT_DIR);
        const pathtostoryfromroot = path.relative(relativepathtoroot, filePath);
        const parsedpath = path.parse(pathtostoryfromroot);
        const csffile = await readCsfFile(pathtostoryfromroot);
        const testcode = await generateTests(pathtostoryfromroot, csffile);
        const testfilename = path.resolve(path.dirname(pathtostoryfromroot), `${parsedpath.name}.test.tsx`);
        await writeTestFile(testfilename, testcode);
        console.groupEnd();
    }
};
function isMDX(filePath) {
    return path.extname(filePath) === EXTENSION_MDX;
}
async function readCsfFile(filePath) {
    console.log(`Reading CSF file ${filePath}`);
    return (await readCsf(filePath, {
        makeTitle: (userTitle) => userTitle || "default",
    })).parse();
}
async function generateTests(filePath, csfFile) {
    const parsedFilePath = path.parse(filePath);
    let code = `
    import { test, expect } from 'vitest'
    import { screen, within } from '@testing-library/react'
    import { composeStories } from '@storybook/react'

    import * as stories from './${parsedFilePath.name}'

    const { ${csfFile.indexInputs.map((indexInput) => indexInput.exportName).join(", ")} } = composeStories(stories)
  `;
    for (const indexInput of csfFile.indexInputs) {
        if (indexInput.type !== "story") {
            continue;
        }
        code += `
    test('${indexInput.name}', async () => {
      await ${indexInput.exportName}.run()
    `;
        code = await autocompleteTest(code, indexInput, csfFile);
    }
    return prettier.format(code, { parser: "typescript" });
}
async function autocompleteTest(code, indexInput, csfFile) {
    const storyExportName = indexInput.exportName;
    const storySource = extractSource(csfFile.getStoryExport(storyExportName));
    const PROMPT_TEMPLATE = await getPromptTemplate();
    const messages = [];
    messages.push(new HumanMessage(Mustache.render(PROMPT_TEMPLATE, { code, storyExportName, storySource })));
    const response = await invokeLlm(messages);
    code += response?.content || "\n})" + "\n";
    return code;
}
async function writeTestFile(fileName, code) {
    console.log("Writing test file", fileName);
    return fs.writeFile(fileName, code, { flag: "w+" });
}
//# sourceMappingURL=test-forge.js.map