jest-roblox-assassin 1.0.0 → 1.1.0

package/src/rewriter.js

@@ -1,32 +1,25 @@
 import chalk from "chalk";
 import fs from "fs";
 import path from "path";
+import util from "util";
 
 export class ResultRewriter {
-    INTERNAL_FRAME_PATTERNS = [
-        /rbxts_include\.node_modules/i,
-        /@rbxts-js\.Promise/i,
-        /@rbxts-js\.JestCircus/i,
-    ];
-
-    constructor({
-        workspaceRoot,
-        projectRoot,
-        rootDir,
-        outDir,
-        datamodelPrefixSegments,
-    }) {
-        this.workspaceRoot = workspaceRoot || projectRoot;
+    constructor({ rojoProject, compilerOptions, testLocationInResults }) {
+        this.rojoProject = rojoProject;
+        this.compilerOptions = compilerOptions;
+        this.testLocationInResults = Boolean(testLocationInResults);
+        this.luauPathMap = new Map();
+        const projectRoot = this.rojoProject?.root ?? process.cwd();
         this.projectRoot = projectRoot;
-        this.rootDir = rootDir;
-        this.outDir = outDir;
-        this.datamodelPrefixSegments = datamodelPrefixSegments;
-
-        const firstSegment = this.datamodelPrefixSegments[0];
-        this.stackFramePattern = new RegExp(
-            `(?:\\[string\\s+")?(${firstSegment}[^\\]":\\n]+)(?:"\\])?:([0-9]+)`,
-            "g"
-        );
+
+        const rootDirRelative = compilerOptions?.rootDir ?? "src";
+        const outDirRelative = compilerOptions?.outDir ?? "out";
+        const absoluteRootDir = path.isAbsolute(rootDirRelative)
+            ? rootDirRelative
+            : path.join(projectRoot, rootDirRelative);
+        const absoluteOutDir = path.isAbsolute(outDirRelative)
+            ? outDirRelative
+            : path.join(projectRoot, outDirRelative);
 
         /**
          * A map from datamodel paths to their corresponding Luau and source file paths.
@@ -34,54 +27,82 @@ export class ResultRewriter {
          */
         this.modulePathMap = (() => {
             const map = new Map();
-            const outRoot = path.join(this.projectRoot, this.outDir);
-            if (!fs.existsSync(outRoot)) return map;
-
-            function visit(folder) {
-                for (const entry of fs.readdirSync(folder, {
-                    withFileTypes: true,
-                })) {
-                    const abs = path.join(folder, entry.name);
-                    if (entry.isDirectory()) {
-                        visit(abs);
-                        continue;
-                    }
-                    if (!entry.name.endsWith(".luau")) continue;
-
-                    const rel = path.relative(outRoot, abs);
-                    const withoutExt = rel.slice(0, -".luau".length);
-                    const datamodelPath = [
-                        ...datamodelPrefixSegments,
-                        ...withoutExt.split(path.sep),
-                    ].join(".");
-
-                    const candidateBases = [
-                        withoutExt + ".ts",
-                        withoutExt + ".tsx",
-                        withoutExt + ".lua",
-                        withoutExt + ".luau",
-                    ];
+            const sourcemap = rojoProject.sourcemap;
+            if (!sourcemap) return map;
+
+            const searchChildren = (node, parents) => {
+                for (const child of node.children) {
+                    // Recurse into children
+                    searchChildren(child, [...parents, child.name]);
+
+                    // Process current child
+                    const datamodelPath = [...parents, child.name].join(".");
+                    const luauPath = child.filePaths[0];
+                    if (!luauPath) continue;
+
+                    const absoluteLuauPath = path.join(projectRoot, luauPath);
                     let sourcePath;
-                    for (const candidate of candidateBases) {
-                        const candidatePath = path.join(
-                            projectRoot,
-                            rootDir,
-                            candidate
+
+                    if (absoluteOutDir) {
+                        const normalizedOutDir = absoluteOutDir.replace(
+                            /\\/g,
+                            "/"
                         );
-                        if (fs.existsSync(candidatePath)) {
-                            sourcePath = candidatePath;
-                            break;
+                        const normalizedLuau = absoluteLuauPath.replace(
+                            /\\/g,
+                            "/"
+                        );
+                        if (normalizedLuau.startsWith(normalizedOutDir)) {
+                            const relativePath = path.relative(
+                                absoluteOutDir,
+                                absoluteLuauPath
+                            );
+                            let candidateSource = path.join(
+                                absoluteRootDir,
+                                relativePath
+                            );
+                            if (
+                                path
+                                    .basename(candidateSource)
+                                    .startsWith("init.")
+                            ) {
+                                candidateSource = path.join(
+                                    path.dirname(candidateSource),
+                                    "index" + path.extname(candidateSource)
+                                );
+                            }
+
+                            if (!fs.existsSync(candidateSource)) {
+                                const withTs = candidateSource.replace(
+                                    /\.(lua|luau)$/,
+                                    ".ts"
+                                );
+                                const withTsx = candidateSource.replace(
+                                    /\.(lua|luau)$/,
+                                    ".tsx"
+                                );
+                                if (fs.existsSync(withTs)) {
+                                    candidateSource = withTs;
+                                } else if (fs.existsSync(withTsx)) {
+                                    candidateSource = withTsx;
+                                }
+                            }
+
+                            if (fs.existsSync(candidateSource)) {
+                                sourcePath = candidateSource;
+                            }
                         }
                     }
 
-                    map.set(datamodelPath, {
-                        luauPath: abs,
-                        sourcePath,
-                    });
+                    const entry = { luauPath, absoluteLuauPath, sourcePath };
+                    map.set(datamodelPath, entry);
+                    const normalizedLuauPath = luauPath.replace(/\\/g, "/");
+                    this.luauPathMap.set(normalizedLuauPath, entry);
                 }
-            }
+            };
+
+            searchChildren(sourcemap, []);
 
-            visit(outRoot);
             return map;
         })();
     }
@@ -131,56 +152,132 @@ export class ResultRewriter {
     }
 
     /**
-     * Formats a file path relative to the project root with forward slashes.
+     * Finds the matcher column for an expectation on a line. Falls back to the start of `expect`.
+     * @param {string} lineText The source line text.
+     * @returns {number} The 1-based column index.
+     */
+    findExpectationColumn(lineText) {
+        if (!lineText) return 1;
+
+        const expectIndex = lineText.search(/\bexpect\s*\(/);
+        if (expectIndex === -1) return 1;
+
+        // Look for the last matcher call after the expect expression to cover any matcher name.
+        const afterExpect = lineText.slice(expectIndex);
+        const matcherRegex = /\.\s*([A-Za-z_$][\w$]*)\s*(?=\()/g;
+        let matcher;
+        let match;
+        while ((match = matcherRegex.exec(afterExpect)) !== null) {
+            matcher = match;
+        }
+
+        if (!matcher) {
+            return expectIndex + 1;
+        }
+
+        const matcherName = matcher[1];
+        const matcherNameOffset =
+            matcher.index + matcher[0].indexOf(matcherName);
+        return expectIndex + matcherNameOffset + 1;
+    }
+
+    /**
+     * Finds the line/column for a test title within a source file.
+     * @param {string} testTitle The title of the test case.
+     * @param {string} sourcePath The path to the source file.
+     * @returns {{ line: number, column: number } | undefined} The location with 1-based line and 0-based column.
+     */
+    findTestHeaderLocation(testTitle, sourcePath) {
+        if (!testTitle || !sourcePath) return undefined;
+        const lines = this.readLines(sourcePath);
+        if (!lines.length) return undefined;
+
+        const escapedTitle = testTitle.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
+        const patterns = [
+            new RegExp(
+                "\\b(?:it|test|xtest|fit|ftest|xit)\\s*\\(\\s*[\"'`]" +
+                    escapedTitle +
+                    "[\"'`]",
+                "i"
+            ),
+            new RegExp("[\"'`]" + escapedTitle + "[\"'`]", "i"),
+        ];
+
+        for (let i = 0; i < lines.length; i++) {
+            const line = lines[i];
+            for (const pattern of patterns) {
+                const match = pattern.exec(line);
+                if (match) {
+                    return { line: i + 1, column: match.index };
+                }
+            }
+        }
+
+        return undefined;
+    }
+
+    /**
+     * Formats a file path ready to be displayed to the terminal by making it relative and using forward slashes.
      * @param {string} filePath The file path to format.
      * @returns {string} The formatted path.
      */
     formatPath(filePath) {
-        return path
-            .relative(this.workspaceRoot, filePath)
-            .split(path.sep)
-            .join("/")
-            .replace(/\\/g, "/");
+        const baseDir = this.projectRoot || process.cwd();
+        let relativePath = path.relative(baseDir, filePath);
+        if (!relativePath || relativePath.startsWith("..")) {
+            relativePath = path.relative(process.cwd(), filePath);
+        }
+        if (!relativePath) {
+            relativePath = filePath;
+        }
+        return relativePath.split(path.sep).join("/").replace(/\\/g, "/");
     }
 
     /**
      * Maps a datamodel stack frame to its source location.
      * @param {string} datamodelPath The path in the datamodel.
      * @param {number} lineNumber The line number in the datamodel file.
-     * @returns {string | undefined} The mapped source location as "path:line" or undefined if not found.
+     * @returns {{ line: number, column: number, file: string } | undefined} The mapped location.
      */
     mapDatamodelFrame(datamodelPath, lineNumber) {
-        const entry = this.modulePathMap.get(datamodelPath);
+        const normalizedPath = datamodelPath.replace(/\\/g, "/");
+        let entry =
+            this.modulePathMap.get(datamodelPath) ||
+            this.modulePathMap.get(normalizedPath);
+        if (!entry) {
+            entry = this.luauPathMap.get(normalizedPath);
+        }
+        if (!entry && this.projectRoot) {
+            const candidateAbsolute = path.isAbsolute(datamodelPath)
+                ? path.resolve(datamodelPath)
+                : path.join(this.projectRoot, normalizedPath);
+            const relativeToRoot = path
+                .relative(this.projectRoot, candidateAbsolute)
+                .replace(/\\/g, "/");
+            entry = this.luauPathMap.get(relativeToRoot);
+        }
         if (!entry) return undefined;
+        const absoluteLuauPath =
+            entry.absoluteLuauPath ?? path.resolve(entry.luauPath);
         const mappedLine = this.findSourceLine(
-            entry.luauPath,
-            entry.sourcePath,
+            absoluteLuauPath,
+            entry.sourcePath ?? absoluteLuauPath,
             lineNumber
         );
-        const displayPath = entry.sourcePath
-            ? this.formatPath(entry.sourcePath)
-            : datamodelPath;
-        return `${displayPath}:${mappedLine}`;
+        const sourceForColumn = entry.sourcePath ?? absoluteLuauPath;
+        const sourceLines = this.readLines(sourceForColumn);
+        const lineText = sourceLines[mappedLine - 1] || "";
+        const column = this.findExpectationColumn(lineText);
+
+        return {
+            line: mappedLine,
+            column,
+            file: sourceForColumn,
+        };
     }
 
     /**
-     * Cleans internal stack trace lines from a text block.
-     * @param {string} text The text block to clean.
-     * @returns {string} The cleaned text block.
-     */
-    cleanInternalLines(text) {
-        if (!text) return text;
-        const lines = text.split(/\r?\n/);
-        const kept = lines.filter(
-            (line) =>
-                !this.INTERNAL_FRAME_PATTERNS.some((pat) => pat.test(line))
-        );
-        const squashed = kept.join("\n").replace(/\n{3,}/g, "\n\n");
-        return squashed.trimEnd();
-    }
-
-    /**
-     * Formats failure messages with syntax highlighting and colors.
+     * Formats `suite.failureMessage` text with colors for terminal output.
      * @param {string} text The failure message text.
      * @returns {string} The formatted failure message.
      */
@@ -234,122 +331,148 @@ export class ResultRewriter {
      * @param {string} value The stack string to rewrite.
      * @returns {string} The rewritten stack string.
      */
-    rewriteStackString(value) {
-        return typeof value === "string"
-            ? this.cleanInternalLines(
-                  value.replace(
-                      this.stackFramePattern,
-                      (match, dmPath, line) => {
-                          const mapped = this.mapDatamodelFrame(
-                              dmPath,
-                              Number(line)
-                          );
-                          return mapped || match;
-                      }
-                  )
-              )
-            : value;
+    rewriteStackString(value, options = {}) {
+        if (!value) return value;
+        const { absolutePaths = false } = options;
+
+        // Try to find matches in modulePathMap
+        const pattern =
+            /((?:[\w@.\/\\-]*[\/.\\][\w@.\/\\-]+)):(\d+)(?::(\d+))?/g;
+        let match;
+        let rewritten = value;
+        const processed = new Set();
+        while ((match = pattern.exec(value)) !== null) {
+            const [fullMatch, filePart, lineStr] = match;
+            const lineNumber = Number(lineStr);
+            if (processed.has(fullMatch)) continue;
+            processed.add(fullMatch);
+            const mapped = this.mapDatamodelFrame(filePart, lineNumber);
+            if (mapped) {
+                const filePath = absolutePaths
+                    ? path.resolve(mapped.file)
+                    : this.formatPath(mapped.file);
+                rewritten = rewritten
+                    .split(fullMatch)
+                    .join(`${filePath}:${mapped.line}:${mapped.column}`);
+            }
+        }
+
+        return rewritten;
     }
 
     /**
-     * Rewrites a test case's failure messages and details.
-     * @param {object} testCase The test case to rewrite.
+     * Rewrites `suite.testResults[].failureMessages` entries.
+     * @param {Array<string>} messages The failure messages array.
+     * @returns {Array<string>} The rewritten messages array.
      */
-    rewriteTestCase(testCase) {
-        if (Array.isArray(testCase.failureMessages)) {
-            const rewritten = testCase.failureMessages.map(
-                this.rewriteStackString.bind(this)
-            );
-            let candidateFrame;
-
-            // Search all failure messages for a valid file frame
-            for (const msg of rewritten) {
-                const frame = this.parseFrame(msg);
-                if (frame) {
-                    candidateFrame = frame;
-                    break;
-                }
-            }
+    rewriteFailureMessages(messages) {
+        if (!Array.isArray(messages)) return messages;
 
-            // Fallback to failureDetails if no frame found in messages
-            if (!candidateFrame && Array.isArray(testCase.failureDetails)) {
-                for (const detail of testCase.failureDetails) {
-                    const stack =
-                        this.rewriteStackString(detail.stack) ||
-                        this.rewriteStackString(detail.__stack) ||
-                        this.rewriteStackString(detail.message);
-                    const frame = this.parseFrame(stack);
-                    if (frame) {
-                        candidateFrame = frame;
-                        break;
-                    }
-                }
-            }
+        const rewriteFailureMessage = (text) => {
+            if (!text) return text;
+
+            const rewritten = this.rewriteStackString(text, {
+                absolutePaths: true,
+            });
 
-            const codeFrame = candidateFrame
-                ? this.buildCodeFrame(
-                      candidateFrame.absPath,
-                      candidateFrame.line,
-                      candidateFrame.column
-                  )
-                : undefined;
-
-            // Only append the code frame to the first message to avoid duplication
-            if (codeFrame && rewritten.length > 0) {
-                rewritten[0] = this.injectCodeFrame(
-                    rewritten[0],
-                    candidateFrame,
-                    codeFrame
+            const lines = rewritten.split(/\r?\n/);
+            const resultLines = [];
+
+            for (const rawLine of lines) {
+                if (!rawLine.trim()) continue;
+
+                // Strip Luau [string "..."] wrappers if present
+                const strippedLine = rawLine.replace(
+                    /^\[string\s+"(.+?)"\]\s*/i,
+                    "$1"
+                );
+
+                const stackMatch = /^(.*):(\d+)(?::(\d+))?(?:\s|$)/.exec(
+                    strippedLine
                 );
+
+                if (stackMatch) {
+                    const [, filePart, lineStr, colStr] = stackMatch;
+                    const lineNumber = Number(lineStr);
+                    const colNumber = Number(colStr || "1");
+                    const mapped = this.mapDatamodelFrame(filePart, lineNumber);
+                    const absFile = mapped
+                        ? path.resolve(mapped.file)
+                        : path.isAbsolute(filePart)
+                        ? filePart
+                        : path.resolve(
+                              this.projectRoot || process.cwd(),
+                              filePart
+                          );
+                    const finalLine = mapped?.line ?? lineNumber;
+                    const finalCol = mapped?.column ?? colNumber;
+                    resultLines.push(
+                        `    at ${absFile}:${finalLine}:${finalCol}`
+                    );
+                    continue;
+                }
+
+                resultLines.push(strippedLine);
             }
 
-            testCase.failureMessages = rewritten.map(
-                this.formatFailureMessage.bind(this)
-            );
-        }
-        if (Array.isArray(testCase.failureDetails)) {
-            testCase.failureDetails = testCase.failureDetails.map((detail) => ({
-                ...detail,
-                stack: this.formatFailureMessage(
-                    this.rewriteStackString(detail.stack)
-                ),
-                __stack: this.formatFailureMessage(
-                    this.rewriteStackString(detail.__stack)
-                ),
-                message: this.formatFailureMessage(
-                    this.rewriteStackString(detail.message)
-                ),
-            }));
-        }
+            return resultLines.join("\n");
+        };
+
+        return messages
+            .map((msg) => rewriteFailureMessage(msg))
+            .filter((msg) => Boolean(msg));
     }
 
     /**
-     * Extends a test file path by checking for common extensions.
-     * @param {string} testFilePath The original test file path.
-     * @returns {string} The extended test file path.
+     * Converts a testFilePath from the Jest datamodel format to the source file path.
+     * A testFilePath looks like: `parent/testName.spec`, where ancestors after parent are sliced off.
+     * Hence, only infering from the modulePathMap is possible. If conflicts arise, no match is made and the original path is returned.
+     *
+     * @param {string} testFilePath The raw testFilePath from runner results.
+     * @returns {string} The source file path.
      */
-    extendTestFilePath(testFilePath) {
+    datamodelPathToSourcePath(testFilePath) {
         if (!testFilePath) return testFilePath;
-        const withExts = [".ts", ".tsx", ""]; // last entry preserves original
-        for (const ext of withExts) {
-            const candidate = path.join(
-                this.projectRoot,
-                `${testFilePath}${ext}`
-            );
-            if (fs.existsSync(candidate)) return candidate;
+
+        // Attempt direct lookup first
+        const normalizedPath = testFilePath.replace(/\\/g, "/");
+        let entry =
+            this.modulePathMap.get(testFilePath) ||
+            this.modulePathMap.get(normalizedPath);
+        if (!entry) {
+            entry = this.luauPathMap.get(normalizedPath);
+        }
+        if (!entry && this.projectRoot) {
+            const candidateAbsolute = path.isAbsolute(testFilePath)
+                ? path.resolve(testFilePath)
+                : path.join(this.projectRoot, normalizedPath);
+            const relativeToRoot = path
+                .relative(this.projectRoot, candidateAbsolute)
+                .replace(/\\/g, "/");
+            entry = this.luauPathMap.get(relativeToRoot);
+        }
+        if (entry?.sourcePath) {
+            return entry.sourcePath;
         }
-        return path.join(this.projectRoot, testFilePath);
-    }
 
-    /**
-     * Strips ANSI escape codes from a string.
-     * @param {string} str The string to strip.
-     * @returns {string} The stripped string.
-     */
-    stripAnsi(str) {
-        return typeof str === "string"
-            ? str.replace(/\u001b\[[0-9;]*m/g, "")
-            : str;
+        let matchingPath = testFilePath
+            .replace(/\.(lua|luau)$/, "")
+            .replace(/\\/g, ".")
+            .replace(/\//g, ".");
+
+        const matches = [];
+        for (const [dmPath, paths] of this.modulePathMap.entries()) {
+            if (dmPath.endsWith(matchingPath)) {
+                matches.push(paths.sourcePath ?? paths.luauPath);
+            }
+        }
+
+        if (matches.length === 1) {
+            return matches[0];
+        }
+
+        // Last resort: return as-is joined with projectRoot
+        return path.join(this.rojoProject.root, testFilePath);
     }
 
     /**
@@ -359,7 +482,7 @@ export class ResultRewriter {
      */
     parseFrame(text) {
         if (!text) return undefined;
-        const cleanText = this.stripAnsi(text);
+        const cleanText = util.stripVTControlCharacters(text);
         // Match patterns like "src/test.ts:10" or "C:\path\test.ts:10"
         // We look for something that looks like a file path followed by a colon and a number
         const pattern =
@@ -370,7 +493,7 @@ export class ResultRewriter {
             const [, filePart, lineStr, colStr] = match;
             const absPath = path.isAbsolute(filePart)
                 ? filePart
-                : path.join(this.workspaceRoot, filePart);
+                : path.join(this.rojoProject.root, filePart);
             if (fs.existsSync(absPath)) {
                 candidates.push({
                     absPath,
@@ -386,10 +509,7 @@ export class ResultRewriter {
         }
 
         if (candidates.length === 0) return undefined;
-        // Sort by score descending, then by line number (prefer higher line numbers which are usually the actual error)
-        return candidates.sort(
-            (a, b) => b.score - a.score || b.line - a.line
-        )[0];
+        return candidates[0];
     }
 
     /**
@@ -421,13 +541,19 @@ export class ResultRewriter {
     buildCodeFrame(absPath, line, column = 1, context = 2) {
         const lines = this.readLines(absPath);
         if (!lines.length) return undefined;
+        // Replace tabs with 4 spaces for consistent formatting
+        lines.forEach((_, idx) => {
+            lines[idx] = lines[idx].replace(/\t/g, "    ");
+        });
         const start = Math.max(1, line - context);
         const end = Math.min(lines.length, line + context + 1);
         const frame = [];
 
+        const digitWidth = String(end).length;
+
         for (let i = start; i <= end; i++) {
             const isBright = i === start;
-            const lineNum = String(i).padStart(String(end).length, " ");
+            const lineNum = String(i).padStart(digitWidth, " ");
             const gutter = `${
                 i === line ? chalk.bold.red(">") : " "
             } ${chalk.grey(lineNum + " |")}`;
@@ -453,7 +579,7 @@ export class ResultRewriter {
     }
 
     /**
-     * Injects a code frame into a text block, moving the corresponding stack frame line to the bottom.
+     * Injects a code frame into a text block, moving all stack trace lines to below the code frame.
      * @param {string} text The text block.
      * @param {object} frame The parsed frame info.
      * @param {string} codeFrame The code frame text.
@@ -463,30 +589,22 @@ export class ResultRewriter {
         if (!codeFrame || !frame) return text;
 
         const lines = text.split(/\r?\n/);
-        let frameLine = "";
-        let frameLineIndex = -1;
-
-        const displayPath = this.formatPath(frame.absPath);
-        const searchPath1 = `${displayPath}:${frame.line}`;
-        const searchPath2 = `${frame.absPath}:${frame.line}`;
+        const stackLines = [];
+        const nonStackLines = [];
+        const pattern =
+            /(?:^|\s|")((?:[a-zA-Z]:[\\\/][^:\s\n"]+|[\w@.\/\\-]+\.[a-z0-9]+)):(\d+)(?::(\d+))?/gi;
 
-        for (let i = 0; i < lines.length; i++) {
-            const line = lines[i];
-            if (line.includes(searchPath1) || line.includes(searchPath2)) {
-                frameLine = line;
-                frameLineIndex = i;
-                break;
+        for (const line of lines) {
+            pattern.lastIndex = 0; // Reset regex state
+            if (pattern.test(line)) {
+                stackLines.push(line);
+            } else {
+                nonStackLines.push(line);
             }
         }
 
-        if (frameLineIndex !== -1) {
-            lines.splice(frameLineIndex, 1);
-            return `${lines
-                .join("\n")
-                .trimEnd()}\n\n${codeFrame}\n\n${frameLine}`;
-        }
-
-        return `${text.trimEnd()}\n\n${codeFrame}`;
+        const mainText = nonStackLines.join("\n");
+        return `${mainText}\n\n${codeFrame}\n\n${stackLines.join("\n")}`;
     }
 
     /**
@@ -495,14 +613,35 @@ export class ResultRewriter {
      */
     rewriteSuiteResult(suite) {
         if (!suite) return;
-        suite.testFilePath = this.formatPath(
-            this.extendTestFilePath(suite.testFilePath)
-        );
+        const sourcePath = this.datamodelPathToSourcePath(suite.testFilePath);
+        const resolvedTestFilePath = path.resolve(sourcePath);
+        suite.testFilePath = resolvedTestFilePath;
 
         if (Array.isArray(suite.testResults)) {
-            suite.testResults.forEach((value) => {
-                this.rewriteTestCase(value);
-            });
+            for (const testResult of suite.testResults) {
+                if (this.testLocationInResults && !testResult.location) {
+                    const location =
+                        this.findTestHeaderLocation(
+                            testResult.title,
+                            sourcePath
+                        ) ||
+                        (sourcePath !== resolvedTestFilePath
+                            ? this.findTestHeaderLocation(
+                                  testResult.title,
+                                  resolvedTestFilePath
+                              )
+                            : undefined);
+                    if (location) {
+                        testResult.location = location;
+                    }
+                }
+
+                if (Array.isArray(testResult.failureMessages)) {
+                    testResult.failureMessages = this.rewriteFailureMessages(
+                        testResult.failureMessages
+                    );
+                }
+            }
         }
 
         if (suite.failureMessage) {
@@ -510,39 +649,29 @@ export class ResultRewriter {
 
             // Split by the test header "  ● " to handle multiple failures in one string
             const sections = rewritten.split(/(\s+●\s+)/);
+
+            const rewriteSection = (sectionContent) => {
+                const frame = this.parseFrame(sectionContent);
+                if (!frame) return sectionContent;
+
+                const codeFrame = this.buildCodeFrame(
+                    frame.absPath,
+                    frame.line,
+                    frame.column
+                );
+                return (
+                    this.injectCodeFrame(sectionContent, frame, codeFrame) +
+                    "\n"
+                );
+            };
+
             if (sections.length > 1) {
                 for (let i = 2; i < sections.length; i += 2) {
-                    const sectionContent = sections[i];
-                    const frame = this.parseFrame(sectionContent);
-                    if (frame) {
-                        const codeFrame = this.buildCodeFrame(
-                            frame.absPath,
-                            frame.line,
-                            frame.column
-                        );
-                        sections[i] =
-                            this.injectCodeFrame(
-                                sectionContent,
-                                frame,
-                                codeFrame
-                            ) + "\n";
-                    }
+                    sections[i] = rewriteSection(sections[i]);
                 }
                 rewritten = sections.join("");
             } else {
-                const frame = this.parseFrame(rewritten);
-                if (frame) {
-                    const codeFrame = this.buildCodeFrame(
-                        frame.absPath,
-                        frame.line,
-                        frame.column
-                    );
-                    rewritten = this.injectCodeFrame(
-                        rewritten,
-                        frame,
-                        codeFrame
-                    ) + "\n";
-                }
+                rewritten = rewriteSection(rewritten);
             }
 
             suite.failureMessage = this.formatFailureMessage(rewritten);
@@ -559,4 +688,72 @@ export class ResultRewriter {
             this.rewriteSuiteResult(suite);
         }
     }
+
+    /**
+     * Rewrites coverage data to use source file paths instead of datamodel paths.
+     * @param {object} coverageData The coverage data.
+     * @returns {object} The rewritten coverage data.
+     */
+    rewriteCoverageData(coverageData) {
+        if (!coverageData) return coverageData;
+
+        const rewritten = {};
+        for (const [datamodelPath, coverage] of Object.entries(coverageData)) {
+            // Skip the "total" key
+            if (datamodelPath === "total") {
+                rewritten[datamodelPath] = coverage;
+                continue;
+            }
+
+            // Convert datamodel path to luau path
+            // Coverage data uses slashes, but modulePathMap uses dots
+            const normalizedPath = datamodelPath.replace(/\//g, ".");
+            let entry = this.modulePathMap.get(normalizedPath);
+
+            const finalPath = path.resolve(entry?.luauPath ?? datamodelPath);
+
+            // Clone the coverage object and update the path property
+            const rewrittenCoverage = { ...coverage };
+            if (rewrittenCoverage.path) {
+                rewrittenCoverage.path = finalPath;
+            }
+
+            rewritten[finalPath] = rewrittenCoverage;
+        }
+
+        return rewritten;
+    }
+
+    /**
+     * Converts raw Jest results into JSON format similar to `jest --json` output.
+     * @param {{ results: object, coverage?: object, globalConfig?: object }} jestRunCliReturn The raw Jest results.
+     */
+    json(jestRunCliReturn) {
+        const results = { ...jestRunCliReturn.results };
+
+        if (jestRunCliReturn.coverage) {
+            results.coverageMap = jestRunCliReturn.coverage;
+        }
+
+        for (const suite of results.testResults) {
+            suite.message = suite.failureMessage ?? "";
+            delete suite.failureMessage;
+
+            suite.assertionResults = suite.testResults || [];
+            delete suite.testResults;
+
+            suite.name = suite.testFilePath;
+            delete suite.testFilePath;
+
+            let overallPassed = true;
+            for (const testResult of suite.testResults || []) {
+                if (testResult.status === "failed") {
+                    overallPassed = false;
+                    break;
+                }
+            }
+            suite.status = overallPassed ? "passed" : "failed";
+        }
+        return results;
+    }
 }