Tests - Fix mountAndGetRootNode usage + Add depth=1 test case

test/traverse-basic.spec.tsx

@@ -1,12 +1,16 @@
-import * as React from "react";
-// import ReactDOM from "react-dom";
-import { mount } from "enzyme";
+// import * as React from "react";
+
+// Import stuff from src
+import { traverse } from "../src";
+
+// Import test helpers and sample components
+import { mountAndGetRootNode } from "./utils/mount-in-enzyme";
+import CDepth1 from "./sample-components/depth-1-simple";
 import CDepth2 from "./sample-components/depth-2-simple";
-import { traverse, FiberNode } from "../src";
 
 describe("Basic traverse test", () => {
   let container: HTMLDivElement;
-
+  
   beforeEach(() => {
     container = document.body.appendChild(document.createElement("div"));
   });
@@ -15,23 +19,23 @@ describe("Basic traverse test", () => {
     document.body.removeChild(container);
   });
 
-  it("should render correctly", () => {
-    const rootRef = React.createRef<CDepth2>();
-
-    // ReactDOM.render(
-    //   <C text="text" ref={rootRef} />,
-    //   document.getElementById("container")
-    // );
-    mount(<CDepth2 ref={rootRef} />, { attachTo: container });
+  it("should work for depth=1", () => {
+    const rootNode = mountAndGetRootNode(CDepth1, container);
 
     const mockCallback = jest.fn();
 
-    if (rootRef.current !== null) {
-      traverse(
-        (rootRef.current as any)._reactInternalFiber.child as FiberNode,
-        mockCallback
-      );
-    }
+    traverse(rootNode, mockCallback);
+
+    // The mock function is called twice
+    expect(mockCallback.mock.calls.length).toBe(1);
+  });
+
+  it("should work for depth=2", () => {
+    const rootNode = mountAndGetRootNode(CDepth2, container);
+
+    const mockCallback = jest.fn();
+
+    traverse(rootNode, mockCallback);
 
     // The mock function is called twice
     expect(mockCallback.mock.calls.length).toBe(2);