comment-to-assert/src/ast-utils.ts

convert line comment to assert.

import * as assert from "node:assert";
import { createRequire } from "node:module";
import {
	type CallExpression,
	type Comment,
	callExpression,
	identifier,
	isCallExpression,
	isDirective,
	isIdentifier,
	isLiteral,
	isNullLiteral,
	stringLiteral,
} from "@babel/types";

const require = createRequire(import.meta.url);
const template = require("@babel/template").default;

const commentCodeRegExp = /=>\s*?(.*?)$/i;

export function tryGetCodeFromComments(comments: ReadonlyArray<Comment>) {
	if (comments.length === 0) {
		return;
	}
	const comment = comments[0];
	if (comment.type === "CommentBlock" || comment.type === "CommentLine") {
		const matchResult = comment.value.match(commentCodeRegExp);
		if (matchResult?.[1]) {
			return matchResult[1];
		}
	}
	return;
}

function isConsole(node: any): node is CallExpression & { expression: any } {
	return isCallExpression(node) && (node.callee as any).object && (node.callee as any).object.name === "console";
}

export const ERROR_COMMENT_PATTERN = /^([a-zA-Z]*?Error)/;
export const PROMISE_RESOLVE_COMMENT_PATTERN = /^Resolve:\s*(.*?)\s*$/;
export const PROMISE_REJECT_COMMENT_PATTERN = /^Reject:\s*(.*?)\s*$/;

/**
 * Creates Babel AST statement(s) from a template string, handling callback placeholders.
 *
 * This function solves the "Found multiple statements but wanted one" error that occurs
 * when Babel template tries to generate multiple statements but expects a single statement.
 * When callbacks are present, it returns an array of statements; otherwise a single statement.
 *
 * @param templateStr - Template string containing BEFORE_CALLBACK and AFTER_CALLBACK placeholders
 * @param templateVars - Variables to substitute in the template
 * @param beforeCallback - Optional callback to execute before the main statement
 * @param afterCallback - Optional callback to execute after the main statement
 * @returns Single statement or array of statements depending on callback presence
 */
function createStatement(templateStr: string, templateVars: any, beforeCallback?: any, afterCallback?: any) {
	if (beforeCallback && afterCallback) {
		const statements = [];
		statements.push(template.statement`BEFORE_CALLBACK;`({ BEFORE_CALLBACK: beforeCallback }));
		const mainTemplate = templateStr.replace(/BEFORE_CALLBACK;|;AFTER_CALLBACK/g, "");
		statements.push(template.statement(mainTemplate)(templateVars));
		statements.push(template.statement`AFTER_CALLBACK;`({ AFTER_CALLBACK: afterCallback }));
		return statements;
	} else if (beforeCallback) {
		const statements = [];
		statements.push(template.statement`BEFORE_CALLBACK;`({ BEFORE_CALLBACK: beforeCallback }));
		const mainTemplate = templateStr.replace(/BEFORE_CALLBACK;|;AFTER_CALLBACK/g, "");
		statements.push(template.statement(mainTemplate)(templateVars));
		return statements;
	} else if (afterCallback) {
		const statements = [];
		const mainTemplate = templateStr.replace(/BEFORE_CALLBACK;|;AFTER_CALLBACK/g, "");
		statements.push(template.statement(mainTemplate)(templateVars));
		statements.push(template.statement`AFTER_CALLBACK;`({ AFTER_CALLBACK: afterCallback }));
		return statements;
	} else {
		// No callbacks, use single statement
		const baseTemplate = templateStr.replace(/BEFORE_CALLBACK;|;AFTER_CALLBACK/g, "");
		return template.statement(baseTemplate)(templateVars);
	}
}

export interface wrapAssertOptions {
	// callback name before assert
	assertBeforeCallbackName?: string;
	// callback name after assert
	assertAfterCallbackName?: string;
}

export function wrapAssert(
	{
		actualNode,
		expectedNode,
		commentExpression,
		id,
	}: { actualNode: any; expectedNode: any; commentExpression: string; id: string },
	options: wrapAssertOptions,
): any {
	assert.notStrictEqual(typeof expectedNode, "undefined");
	const ACTUAL_NODE = actualNode;
	const EXPECTED_NODE = expectedNode;
	const BEFORE_CALLBACK = options.assertBeforeCallbackName
		? callExpression(identifier(options.assertBeforeCallbackName), [stringLiteral(id)])
		: undefined;
	const AFTER_CALLBACK = options.assertAfterCallbackName
		? callExpression(identifier(options.assertAfterCallbackName), [stringLiteral(id)])
		: undefined;
	if (isConsole(actualNode)) {
		const args = actualNode.arguments;
		const firstArgument = args[0];
		return wrapAssert(
			{
				actualNode: firstArgument,
				expectedNode,
				commentExpression,
				id,
			},
			options,
		);
	} else if (isIdentifier(expectedNode) && ERROR_COMMENT_PATTERN.test(expectedNode.name)) {
		return createStatement(
			"BEFORE_CALLBACK;assert.throws(function() { ACTUAL_NODE });AFTER_CALLBACK;",
			{ ACTUAL_NODE },
			BEFORE_CALLBACK,
			AFTER_CALLBACK,
		);
	} else if (expectedNode.type === "Resolve") {
		// getExpressionNodeFromCommentValue define the type
		const ARGS = isConsole(actualNode) ? actualNode.arguments[0] : actualNode;
		return template.statement`Promise.resolve(ARGS).then(v => {
            ${wrapAssert(
				{
					actualNode: { type: "Identifier", name: "v" },
					expectedNode: expectedNode.node,
					commentExpression,
					id,
				},
				options,
			)}
            return v;
        });`({
			ARGS,
		});
	} else if (expectedNode.type === "Reject") {
		const ARGS = isConsole(actualNode) ? actualNode.arguments[0] : actualNode;
		if (BEFORE_CALLBACK || AFTER_CALLBACK) {
			if (AFTER_CALLBACK) {
				return createStatement(
					"BEFORE_CALLBACK;assert.rejects(ARGS).then(() => { AFTER_CALLBACK; });",
					{ ARGS, AFTER_CALLBACK },
					BEFORE_CALLBACK,
					undefined,
				);
			} else {
				return createStatement(
					"BEFORE_CALLBACK;assert.rejects(ARGS);AFTER_CALLBACK;",
					{ ARGS },
					BEFORE_CALLBACK,
					AFTER_CALLBACK,
				);
			}
		} else {
			return template.statement`assert.rejects(ARGS).then(() => {});`({
				ARGS,
			});
		}
	} else if (isIdentifier(expectedNode) && expectedNode.name === "NaN") {
		return createStatement(
			"BEFORE_CALLBACK;assert.ok(isNaN(ACTUAL_NODE));AFTER_CALLBACK;",
			{ ACTUAL_NODE },
			BEFORE_CALLBACK,
			AFTER_CALLBACK,
		);
	} else if (isNullLiteral(expectedNode)) {
		return createStatement(
			"BEFORE_CALLBACK;assert.strictEqual(ACTUAL_NODE, null);AFTER_CALLBACK;",
			{ ACTUAL_NODE },
			BEFORE_CALLBACK,
			AFTER_CALLBACK,
		);
	} else if (isIdentifier(expectedNode) && expectedNode.name === "undefined") {
		return createStatement(
			"BEFORE_CALLBACK;assert.strictEqual(ACTUAL_NODE, undefined);AFTER_CALLBACK;",
			{ ACTUAL_NODE },
			BEFORE_CALLBACK,
			AFTER_CALLBACK,
		);
	} else if (isLiteral(expectedNode)) {
		// Handle Directive Prorogue as string literal
		if (isDirective(ACTUAL_NODE)) {
			const actualValue = (ACTUAL_NODE.value as any).extra.raw;
			return createStatement(
				"BEFORE_CALLBACK;assert.strictEqual(ACTUAL_NODE, EXPECTED_NODE);AFTER_CALLBACK;",
				{ ACTUAL_NODE: actualValue, EXPECTED_NODE },
				BEFORE_CALLBACK,
				AFTER_CALLBACK,
			);
		} else {
			return createStatement(
				"BEFORE_CALLBACK;assert.strictEqual(ACTUAL_NODE, EXPECTED_NODE);AFTER_CALLBACK;",
				{ ACTUAL_NODE, EXPECTED_NODE },
				BEFORE_CALLBACK,
				AFTER_CALLBACK,
			);
		}
	} else {
		return createStatement(
			"BEFORE_CALLBACK;assert.deepStrictEqual(ACTUAL_NODE, EXPECTED_NODE);AFTER_CALLBACK;",
			{ ACTUAL_NODE, EXPECTED_NODE },
			BEFORE_CALLBACK,
			AFTER_CALLBACK,
		);
	}
	throw new Error(`Unknown pattern: ${actualNode}`);
}