prefer-find-by.ts

import { TSESTree, ASTUtils, TSESLint } from '@typescript-eslint/utils';

import { createTestingLibraryRule } from '../create-testing-library-rule';
import {
	isArrowFunctionExpression,
	isCallExpression,
	isMemberExpression,
	isObjectExpression,
	isObjectPattern,
	isProperty,
} from '../node-utils';

export const RULE_NAME = 'prefer-find-by';
export type MessageIds = 'preferFindBy';
type Options = [];

export const WAIT_METHODS = ['waitFor', 'waitForElement', 'wait'] as const;

export function getFindByQueryVariant(
	queryMethod: string
): 'findAllBy' | 'findBy' {
	return queryMethod.includes('All') ? 'findAllBy' : 'findBy';
}

function findRenderDefinitionDeclaration(
	scope: TSESLint.Scope.Scope | null,
	query: string
): TSESTree.Identifier | null {
	if (!scope) {
		return null;
	}

	const variable = scope.variables.find(
		(v: TSESLint.Scope.Variable) => v.name === query
	);

	if (variable) {
		return (
			variable.defs
				.map(({ name }) => name)
				.filter(ASTUtils.isIdentifier)
				.find(({ name }) => name === query) ?? null
		);
	}

	return findRenderDefinitionDeclaration(scope.upper, query);
}

export default createTestingLibraryRule<Options, MessageIds>({
	name: RULE_NAME,
	meta: {
		type: 'suggestion',
		docs: {
			description:
				'Suggest using `find(All)By*` query instead of `waitFor` + `get(All)By*` to wait for elements',
			recommendedConfig: {
				dom: 'error',
				angular: 'error',
				react: 'error',
				vue: 'error',
				marko: 'error',
			},
		},
		messages: {
			preferFindBy:
				'Prefer `{{queryVariant}}{{queryMethod}}` query over using `{{waitForMethodName}}` + `{{prevQuery}}`',
		},
		fixable: 'code',
		schema: [],
	},
	defaultOptions: [],

	create(context, _, helpers) {
		const sourceCode = context.getSourceCode();

		/**
		 * Reports the invalid usage of wait* plus getBy/QueryBy methods and automatically fixes the scenario
		 * @param node - The CallExpresion node that contains the wait* method
		 * @param replacementParams - Object with info for error message and autofix:
		 * @param replacementParams.queryVariant - The variant method used to query: findBy/findAllBy.
		 * @param replacementParams.prevQuery - The query originally used inside `waitFor`
		 * @param replacementParams.queryMethod - Suffix string to build the query method (the query-part that comes after the "By"): LabelText, Placeholder, Text, Role, Title, etc.
		 * @param replacementParams.waitForMethodName - wait for method used: waitFor/wait/waitForElement
		 * @param replacementParams.fix - Function that applies the fix to correct the code
		 */
		function reportInvalidUsage(
			node: TSESTree.CallExpression,
			replacementParams: {
				queryVariant: 'findAllBy' | 'findBy';
				queryMethod: string;
				prevQuery: string;
				waitForMethodName: string;
				fix: TSESLint.ReportFixFunction;
			}
		) {
			const { queryMethod, queryVariant, prevQuery, waitForMethodName, fix } =
				replacementParams;
			context.report({
				node,
				messageId: 'preferFindBy',
				data: {
					queryVariant,
					queryMethod,
					prevQuery,
					waitForMethodName,
				},
				fix,
			});
		}

		function getWrongQueryNameInAssertion(
			node: TSESTree.ArrowFunctionExpression
		) {
			if (
				!isCallExpression(node.body) ||
				!isMemberExpression(node.body.callee)
			) {
				return null;
			}

			// expect(getByText).toBeInTheDocument() shape
			if (
				isCallExpression(node.body.callee.object) &&
				isCallExpression(node.body.callee.object.arguments[0]) &&
				ASTUtils.isIdentifier(node.body.callee.object.arguments[0].callee)
			) {
				return node.body.callee.object.arguments[0].callee.name;
			}

			if (!ASTUtils.isIdentifier(node.body.callee.property)) {
				return null;
			}

			// expect(screen.getByText).toBeInTheDocument() shape
			if (
				isCallExpression(node.body.callee.object) &&
				isCallExpression(node.body.callee.object.arguments[0]) &&
				isMemberExpression(node.body.callee.object.arguments[0].callee) &&
				ASTUtils.isIdentifier(
					node.body.callee.object.arguments[0].callee.property
				)
			) {
				return node.body.callee.object.arguments[0].callee.property.name;
			}

			// expect(screen.getByText).not shape
			if (
				isMemberExpression(node.body.callee.object) &&
				isCallExpression(node.body.callee.object.object) &&
				isCallExpression(node.body.callee.object.object.arguments[0]) &&
				isMemberExpression(
					node.body.callee.object.object.arguments[0].callee
				) &&
				ASTUtils.isIdentifier(
					node.body.callee.object.object.arguments[0].callee.property
				)
			) {
				return node.body.callee.object.object.arguments[0].callee.property.name;
			}

			// expect(getByText).not shape
			if (
				isMemberExpression(node.body.callee.object) &&
				isCallExpression(node.body.callee.object.object) &&
				isCallExpression(node.body.callee.object.object.arguments[0]) &&
				ASTUtils.isIdentifier(
					node.body.callee.object.object.arguments[0].callee
				)
			) {
				return node.body.callee.object.object.arguments[0].callee.name;
			}

			return node.body.callee.property.name;
		}

		function getWrongQueryName(node: TSESTree.ArrowFunctionExpression) {
			if (!isCallExpression(node.body)) {
				return null;
			}

			// expect(() => getByText) and expect(() => screen.getByText) shape
			if (
				ASTUtils.isIdentifier(node.body.callee) &&
				helpers.isSyncQuery(node.body.callee)
			) {
				return node.body.callee.name;
			}

			return getWrongQueryNameInAssertion(node);
		}

		function getCaller(node: TSESTree.ArrowFunctionExpression) {
			if (
				!isCallExpression(node.body) ||
				!isMemberExpression(node.body.callee)
			) {
				return null;
			}

			if (ASTUtils.isIdentifier(node.body.callee.object)) {
				// () => screen.getByText
				return node.body.callee.object.name;
			}

			if (
				// expect()
				isCallExpression(node.body.callee.object) &&
				ASTUtils.isIdentifier(node.body.callee.object.callee) &&
				isCallExpression(node.body.callee.object.arguments[0]) &&
				isMemberExpression(node.body.callee.object.arguments[0].callee) &&
				ASTUtils.isIdentifier(
					node.body.callee.object.arguments[0].callee.object
				)
			) {
				return node.body.callee.object.arguments[0].callee.object.name;
			}

			if (
				// expect().not
				isMemberExpression(node.body.callee.object) &&
				isCallExpression(node.body.callee.object.object) &&
				isCallExpression(node.body.callee.object.object.arguments[0]) &&
				isMemberExpression(
					node.body.callee.object.object.arguments[0].callee
				) &&
				ASTUtils.isIdentifier(
					node.body.callee.object.object.arguments[0].callee.object
				)
			) {
				return node.body.callee.object.object.arguments[0].callee.object.name;
			}

			return null;
		}

		function isSyncQuery(node: TSESTree.ArrowFunctionExpression) {
			if (!isCallExpression(node.body)) {
				return false;
			}

			const isQuery =
				ASTUtils.isIdentifier(node.body.callee) &&
				helpers.isSyncQuery(node.body.callee);

			const isWrappedInPresenceAssert =
				isMemberExpression(node.body.callee) &&
				isCallExpression(node.body.callee.object) &&
				isCallExpression(node.body.callee.object.arguments[0]) &&
				ASTUtils.isIdentifier(node.body.callee.object.arguments[0].callee) &&
				helpers.isSyncQuery(node.body.callee.object.arguments[0].callee) &&
				helpers.isPresenceAssert(node.body.callee);

			const isWrappedInNegatedPresenceAssert =
				isMemberExpression(node.body.callee) &&
				isMemberExpression(node.body.callee.object) &&
				isCallExpression(node.body.callee.object.object) &&
				isCallExpression(node.body.callee.object.object.arguments[0]) &&
				ASTUtils.isIdentifier(
					node.body.callee.object.object.arguments[0].callee
				) &&
				helpers.isSyncQuery(
					node.body.callee.object.object.arguments[0].callee
				) &&
				helpers.isPresenceAssert(node.body.callee.object);

			return (
				isQuery || isWrappedInPresenceAssert || isWrappedInNegatedPresenceAssert
			);
		}

		function isScreenSyncQuery(node: TSESTree.ArrowFunctionExpression) {
			if (!isArrowFunctionExpression(node) || !isCallExpression(node.body)) {
				return false;
			}

			if (
				!isMemberExpression(node.body.callee) ||
				!ASTUtils.isIdentifier(node.body.callee.property)
			) {
				return false;
			}

			if (
				!ASTUtils.isIdentifier(node.body.callee.object) &&
				!isCallExpression(node.body.callee.object) &&
				!isMemberExpression(node.body.callee.object)
			) {
				return false;
			}

			const isWrappedInPresenceAssert =
				helpers.isPresenceAssert(node.body.callee) &&
				isCallExpression(node.body.callee.object) &&
				isCallExpression(node.body.callee.object.arguments[0]) &&
				isMemberExpression(node.body.callee.object.arguments[0].callee) &&
				ASTUtils.isIdentifier(
					node.body.callee.object.arguments[0].callee.object
				);

			const isWrappedInNegatedPresenceAssert =
				isMemberExpression(node.body.callee.object) &&
				helpers.isPresenceAssert(node.body.callee.object) &&
				isCallExpression(node.body.callee.object.object) &&
				isCallExpression(node.body.callee.object.object.arguments[0]) &&
				isMemberExpression(node.body.callee.object.object.arguments[0].callee);

			return (
				helpers.isSyncQuery(node.body.callee.property) ||
				isWrappedInPresenceAssert ||
				isWrappedInNegatedPresenceAssert
			);
		}

		function getQueryArguments(node: TSESTree.CallExpression) {
			if (
				isMemberExpression(node.callee) &&
				isCallExpression(node.callee.object) &&
				isCallExpression(node.callee.object.arguments[0])
			) {
				return node.callee.object.arguments[0].arguments;
			}

			if (
				isMemberExpression(node.callee) &&
				isMemberExpression(node.callee.object) &&
				isCallExpression(node.callee.object.object) &&
				isCallExpression(node.callee.object.object.arguments[0])
			) {
				return node.callee.object.object.arguments[0].arguments;
			}

			return node.arguments;
		}

		return {
			'AwaitExpression > CallExpression'(node: TSESTree.CallExpression) {
				if (
					!ASTUtils.isIdentifier(node.callee) ||
					!helpers.isAsyncUtil(node.callee, WAIT_METHODS)
				) {
					return;
				}
				// ensure the only argument is an arrow function expression - if the arrow function is a block
				// we skip it
				const argument = node.arguments[0];
				if (
					!isArrowFunctionExpression(argument) ||
					!isCallExpression(argument.body)
				) {
					return;
				}

				const waitForMethodName = node.callee.name;

				// ensure here it's one of the sync methods that we are calling
				if (isScreenSyncQuery(argument)) {
					const caller = getCaller(argument);

					if (!caller) {
						return;
					}

					// shape of () => screen.getByText
					const fullQueryMethod = getWrongQueryName(argument);

					if (!fullQueryMethod) {
						return;
					}

					// if there is a second argument to AwaitExpression, it is the options
					const waitOptions = node.arguments[1];
					let waitOptionsSourceCode = '';
					if (isObjectExpression(waitOptions)) {
						waitOptionsSourceCode = `, ${sourceCode.getText(waitOptions)}`;
					}

					const queryVariant = getFindByQueryVariant(fullQueryMethod);
					const callArguments = getQueryArguments(argument.body);
					const queryMethod = fullQueryMethod.split('By')[1];

					if (!queryMethod) {
						return;
					}

					reportInvalidUsage(node, {
						queryMethod,
						queryVariant,
						prevQuery: fullQueryMethod,
						waitForMethodName,
						fix(fixer) {
							const property = (
								(argument.body as TSESTree.CallExpression)
									.callee as TSESTree.MemberExpression
							).property;
							if (helpers.isCustomQuery(property as TSESTree.Identifier)) {
								return null;
							}
							const newCode = `${caller}.${queryVariant}${queryMethod}(${callArguments
								.map((callArgNode) => sourceCode.getText(callArgNode))
								.join(', ')}${waitOptionsSourceCode})`;
							return fixer.replaceText(node, newCode);
						},
					});
					return;
				}

				if (!isSyncQuery(argument)) {
					return;
				}

				// shape of () => getByText
				const fullQueryMethod = getWrongQueryName(argument);

				if (!fullQueryMethod) {
					return;
				}

				const queryMethod = fullQueryMethod.split('By')[1];
				const queryVariant = getFindByQueryVariant(fullQueryMethod);
				const callArguments = getQueryArguments(argument.body);

				reportInvalidUsage(node, {
					queryMethod,
					queryVariant,
					prevQuery: fullQueryMethod,
					waitForMethodName,
					fix(fixer) {
						// we know from above callee is an Identifier
						if (
							helpers.isCustomQuery(
								(argument.body as TSESTree.CallExpression)
									.callee as TSESTree.Identifier
							)
						) {
							return null;
						}
						const findByMethod = `${queryVariant}${queryMethod}`;
						const allFixes: TSESLint.RuleFix[] = [];
						// this updates waitFor with findBy*
						const newCode = `${findByMethod}(${callArguments
							.map((callArgNode) => sourceCode.getText(callArgNode))
							.join(', ')})`;
						allFixes.push(fixer.replaceText(node, newCode));

						// this adds the findBy* declaration - adding it to the list of destructured variables { findBy* } = render()
						const definition = findRenderDefinitionDeclaration(
							context.getScope(),
							fullQueryMethod
						);
						// I think it should always find it, otherwise code should not be valid (it'd be using undeclared variables)
						if (!definition) {
							return allFixes;
						}
						// check the declaration is part of a destructuring
						if (
							definition.parent &&
							isObjectPattern(definition.parent.parent)
						) {
							const allVariableDeclarations = definition.parent.parent;
							// verify if the findBy* method was already declared
							if (
								allVariableDeclarations.properties.some(
									(p) =>
										isProperty(p) &&
										ASTUtils.isIdentifier(p.key) &&
										p.key.name === findByMethod
								)
							) {
								return allFixes;
							}
							// the last character of a destructuring is always a  "}", so we should replace it with the findBy* declaration
							const textDestructuring = sourceCode.getText(
								allVariableDeclarations
							);
							const text = textDestructuring.replace(
								/(\s*})$/,
								`, ${findByMethod}$1`
							);
							allFixes.push(fixer.replaceText(allVariableDeclarations, text));
						}

						return allFixes;
					},
				});
			},
		};
	},
});