react-fiber-traverse/test/utils/mountInEnzyme.tsx

import * as React from "react";
import { mount } from "enzyme";
import {
  FiberNode,
  FiberNodeForComponentClass,
} from "../../src/mocked-types";
import {
  isConstructorComponentClass
} from "../../src/utils";
import getWrappedComponent from "./getWrappedComponent";

export class RootNodeNotFoundError extends Error {
  static message = `Couldn't find root node. This might occur if render has become async`;
  constructor() {
    super(RootNodeNotFoundError.message);
    this.name = "RootNodeNotFoundError";
  }
}

export function mountAndGetRootNode(
  SomeComponent: React.ElementType,
  container: HTMLElement
): FiberNode {
  const rootRef = React.createRef<React.Component>();

  let ToMount: React.ComponentClass | undefined = undefined;
  // Wrap function components in ComponentClass to be able to set ref
  if (!isConstructorComponentClass(SomeComponent)) {
    ToMount = getWrappedComponent(SomeComponent);
  } else {
    ToMount = SomeComponent;
  }

  // TODO: later make this async as mount might be async and hence next steps will fail
  mount(<ToMount ref={rootRef} />, { attachTo: container });

  if (rootRef.current === null) {
    throw new RootNodeNotFoundError();
  }

  const rootNode = ((rootRef.current as any)._reactInternalFiber as FiberNode)
    .child;

  if (rootNode === null) {
    throw new RootNodeNotFoundError();
  }

  // Unwrap function components to return expected root node
  if (!isConstructorComponentClass(SomeComponent)) {
    return (rootNode as FiberNodeForComponentClass).child as FiberNode;
  } else {
    return rootNode;
  }
}