diff --git a/test/sample-components/depth-1-simple.tsx b/test/sample-components/depth-1-simple.tsx
new file mode 100644
index 0000000..18cb464
--- /dev/null
+++ b/test/sample-components/depth-1-simple.tsx
@@ -0,0 +1,10 @@
+// import * as PropTypes from "prop-types";
+import * as React from "react";
+
+class CDepth1 extends React.Component {
+  render() {
+    return null;
+  }
+}
+
+export default CDepth1;
diff --git a/test/sample-components/depth-2-simple.tsx b/test/sample-components/depth-2-simple.tsx
index b6a0e66..3788cdf 100644
--- a/test/sample-components/depth-2-simple.tsx
+++ b/test/sample-components/depth-2-simple.tsx
@@ -7,8 +7,4 @@ class CDepth2 extends React.Component {
   }
 }
 
-// C.propTypes = {
-//   text: PropTypes.string.isRequired
-// };
-
 export default CDepth2;
diff --git a/test/traverse-basic.spec.tsx b/test/traverse-basic.spec.tsx
index cd02533..16cfa75 100644
--- a/test/traverse-basic.spec.tsx
+++ b/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);