test.jsx

import { isSettledRepeatNode, SettledRootElementNodeMap } from '@/core';
import React, { ReactInstance, useLayoutEffect } from 'react';
import { getComponentEnhancer } from '../constants';

export const TEST_ID = 'data-testid';

export function isDOMNode(node: unknown): node is Element | Text {
  const t = node ? (node as { nodeType: unknown }).nodeType : null;
  return t ? t === Node.ELEMENT_NODE || t === Node.TEXT_NODE : false;
}

export interface FiberNode {
  return?: FiberNode | null;
  stateNode?: ReactInstance;
  sibling?: FiberNode;
  child?: FiberNode;
  key?: string;
}

export const elementsFromFiber = (fiber: FiberNode, elements: Array<Element | Text>): void => {
  if (fiber.stateNode && isDOMNode(fiber.stateNode)) {
    elements.push(fiber.stateNode);
  } else if (fiber.child) {
    elementsFromFiber(fiber.child, elements);
  }

  if (fiber.sibling) {
    elementsFromFiber(fiber.sibling, elements);
  }
};

export const setRuntimeId = ({
  elements,
  isRepeatNode,
  fiberNode,
  currentElementId
}: {
  currentElementId: string;
  elements: Element[];
  isRepeatNode: boolean;
  fiberNode: FiberNode;
}): void => {
  const testId = isRepeatNode ? `${currentElementId}_repeat(${fiberNode.key})` : `${currentElementId}`;

  if (elements.length > 1) {
    elements.forEach((currentElement, index) => {
      currentElement.setAttribute(TEST_ID, `${testId}_fragment(${index})`);
    });
  } else {
    const currentElement = elements[0];

    currentElement.setAttribute(TEST_ID, testId);
  }
};

export const injectTestId = getComponentEnhancer(
  ({ component, settledElement }, { app }): React.FC => {
    return (props): JSX.Element => {
      const Comp = component;
      const element = <Comp {...props} />;

      useLayoutEffect(() => {
        const elements: Element[] = [];

        const { $rootId, id } = settledElement;
        if (!id) {
          return;
        }

        const settledNode = (
          app as unknown as {
            getSettledRootElementNode: (rootId: string) => SettledRootElementNodeMap | undefined;
          }
        )
          .getSettledRootElementNode($rootId as string)
          ?.get(id);
        const isRepeatNode = isSettledRepeatNode(settledNode);

        const fiberNode = (element as unknown as { _owner: FiberNode })._owner;
        if (fiberNode?.child) {
          elementsFromFiber(fiberNode, elements);
        }

        if (elements.length) {
          setRuntimeId({
            elements,
            currentElementId: settledElement.id,
            fiberNode,
            isRepeatNode
          });
        }
        // eslint-disable-next-line react-hooks/exhaustive-deps
      }, []);

      return element;
    };
  },
  1,
  'inject-test-id',
  false
);