feat(RTTR): handle primitives in test-renderer and fix queries in TestInstances (#3507)

Author: AlaricBaraou

jest.config.js

@@ -10,7 +10,6 @@ module.exports = {
     '<rootDir>/packages/fiber/dist',
     '<rootDir>/packages/fiber/src/index',
     '<rootDir>/packages/test-renderer/dist',
-    '<rootDir>/test-utils',
   ],
   coverageDirectory: './coverage/',
   collectCoverage: false,

packages/test-renderer/markdown/rttr.md

@@ -125,7 +125,7 @@ Similar to the [`act()` in `react-test-renderer`](https://reactjs.org/docs/test-
 #### Act example (using jest)
 
 ```tsx
-import ReactThreeTestRenderer from 'react-three-test-renderer'
+import ReactThreeTestRenderer from '@react-three/test-renderer'
 
 const Mesh = () => {
   const meshRef = React.useRef()

packages/test-renderer/src/__tests__/RTTR.core.test.tsx

@@ -351,10 +351,100 @@ describe('ReactThreeTestRenderer Core', () => {
     expect(renderer.toTree()).toMatchSnapshot()
   })
 
+  it('correctly searches through multiple levels in regular objects', async () => {
+    // Create a deep tree: group -> mesh -> mesh -> mesh
+    const renderer = await ReactThreeTestRenderer.create(
+      <group name="root-group">
+        <mesh name="level1-mesh">
+          <boxGeometry />
+          <meshBasicMaterial color="red" />
+          <mesh name="level2-mesh">
+            <boxGeometry />
+            <meshBasicMaterial color="green" />
+            <mesh name="level3-mesh">
+              <boxGeometry />
+              <meshBasicMaterial color="blue" />
+            </mesh>
+          </mesh>
+        </mesh>
+      </group>,
+    )
+
+    // Test from the root
+    const allMeshes = renderer.scene.findAllByType('Mesh')
+    expect(allMeshes.length).toBe(3) // Should find all three meshes
+
+    // Test from an intermediate node
+    const topMesh = renderer.scene.find((node) => node.props.name === 'level1-mesh')
+    const nestedMeshes = topMesh.findAllByType('Mesh')
+    expect(nestedMeshes.length).toBe(2) // Should find the two nested meshes
+
+    // Find a deeply nested mesh from an intermediate node by property
+    const level3 = topMesh.find((node) => node.props.name === 'level3-mesh')
+    expect(level3).toBeDefined()
+    expect(level3.type).toBe('Mesh')
+  })
+
+  it('Can search from retrieved primitive Instance', async () => {
+    const group = new THREE.Group()
+    group.name = 'PrimitiveGroup'
+
+    const childMesh = new THREE.Mesh(new THREE.BoxGeometry(1, 1, 1), new THREE.MeshBasicMaterial({ color: 'red' }))
+    childMesh.name = 'PrimitiveChildMesh'
+    group.add(childMesh)
+
+    const nestedMesh = new THREE.Mesh(new THREE.BoxGeometry(1, 1, 1), new THREE.MeshBasicMaterial({ color: 'red' }))
+    nestedMesh.name = 'PrimitiveNestedChildMesh'
+    childMesh.add(nestedMesh)
+
+    const renderer = await ReactThreeTestRenderer.create(<primitive object={group} />)
+
+    const foundGroup = renderer.scene.findByType('Group')
+    const foundMesh = foundGroup.children[0]
+    const foundNestedMesh = foundMesh.findByType('Mesh')
+    expect(foundNestedMesh).toBeDefined()
+  })
+
   it('root instance and refs return the same value', async () => {
     let refInst = null
     const renderer = await ReactThreeTestRenderer.create(<mesh ref={(ref) => (refInst = ref)} />)
     const root = renderer.getInstance() // this will be Mesh
     expect(root).toEqual(refInst)
   })
+
+  it('handles primitive objects and their children correctly in toGraph', async () => {
+    // Create a component with both regular objects and primitives with children
+    const PrimitiveTestComponent = () => {
+      // Create a THREE.js group with mesh children
+      const group = new THREE.Group()
+      group.name = 'PrimitiveGroup'
+
+      const childMesh = new THREE.Mesh(new THREE.BoxGeometry(1, 1, 1), new THREE.MeshBasicMaterial({ color: 'red' }))
+      childMesh.name = 'PrimitiveChildMesh'
+      group.add(childMesh)
+
+      // Add a nested group to test deeper hierarchies
+      const nestedGroup = new THREE.Group()
+      nestedGroup.name = 'NestedGroup'
+      const nestedMesh = new THREE.Mesh(new THREE.SphereGeometry(0.5), new THREE.MeshBasicMaterial({ color: 'blue' }))
+      nestedMesh.name = 'NestedMesh'
+      nestedGroup.add(nestedMesh)
+      group.add(nestedGroup)
+
+      return (
+        <>
+          <mesh name="RegularMesh">
+            <boxGeometry args={[2, 2]} />
+            <meshBasicMaterial />
+          </mesh>
+
+          <primitive object={group} />
+        </>
+      )
+    }
+
+    const renderer = await ReactThreeTestRenderer.create(<PrimitiveTestComponent />)
+
+    expect(renderer.toGraph()).toMatchSnapshot()
+  })
 })

packages/test-renderer/src/__tests__/__snapshots__/RTTR.core.test.tsx.snap

@@ -171,6 +171,49 @@ Array [
 ]
 `;
 
+exports[`ReactThreeTestRenderer Core handles primitive objects and their children correctly in toGraph 1`] = `
+Array [
+  Object {
+    "children": Array [
+      Object {
+        "children": Array [],
+        "name": "",
+        "type": "BoxGeometry",
+      },
+      Object {
+        "children": Array [],
+        "name": "",
+        "type": "MeshBasicMaterial",
+      },
+    ],
+    "name": "RegularMesh",
+    "type": "Mesh",
+  },
+  Object {
+    "children": Array [
+      Object {
+        "children": Array [],
+        "name": "PrimitiveChildMesh",
+        "type": "Mesh",
+      },
+      Object {
+        "children": Array [
+          Object {
+            "children": Array [],
+            "name": "NestedMesh",
+            "type": "Mesh",
+          },
+        ],
+        "name": "NestedGroup",
+        "type": "Group",
+      },
+    ],
+    "name": "PrimitiveGroup",
+    "type": "Group",
+  },
+]
+`;
+
 exports[`ReactThreeTestRenderer Core toTree() handles complicated tree of fragments 1`] = `
 Array [
   Object {

packages/test-renderer/src/createTestInstance.ts

@@ -5,6 +5,31 @@ import type { Obj, TestInstanceChildOpts } from './types/internal'
 
 import { expectOne, matchProps, findAll } from './helpers/testInstance'
 
+// Helper to create a minimal wrapper for THREE.Object3D children of primitives
+const createVirtualInstance = (object: THREE.Object3D, parent: Instance<any>): Instance<THREE.Object3D> => {
+  // Create the virtual instance for this object
+  // we can't import the prepare method from packages/fiber/src/core/utils.tsx so we do what we can
+  const instance: Instance<THREE.Object3D> = {
+    root: parent.root,
+    type: object.type.toLowerCase(), // Convert to lowercase to match R3F convention
+    parent,
+    children: [],
+    props: { object },
+    object,
+    eventCount: 0,
+    handlers: {},
+    isHidden: false,
+  }
+
+  // Recursively process children if they exist
+  if (object.children && object.children.length > 0) {
+    const objectChildren = object.children as THREE.Object3D[]
+    instance.children = Array.from(objectChildren).map((child) => createVirtualInstance(child, instance))
+  }
+
+  return instance
+}
+
 export class ReactThreeTestInstance<TObject extends THREE.Object3D = THREE.Object3D> {
   _fiber: Instance<TObject>
 
@@ -47,11 +72,30 @@ export class ReactThreeTestInstance<TObject extends THREE.Object3D = THREE.Objec
   private getChildren = (
     fiber: Instance,
     opts: TestInstanceChildOpts = { exhaustive: false },
-  ): ReactThreeTestInstance[] =>
-    fiber.children.filter((child) => !child.props.attach || opts.exhaustive).map((fib) => wrapFiber(fib as Instance))
+  ): ReactThreeTestInstance[] => {
+    // Get standard R3F children
+    const r3fChildren = fiber.children
+      .filter((child) => !child.props.attach || opts.exhaustive)
+      .map((fib) => wrapFiber(fib as Instance))
+
+    // For primitives, also add THREE.js object children
+    if (fiber.type === 'primitive' && fiber.object.children?.length) {
+      const threeChildren = fiber.object.children.map((child: THREE.Object3D) => {
+        // Create a virtual instance that wraps the THREE.js child
+        const virtualInstance = createVirtualInstance(child, fiber)
+        return wrapFiber(virtualInstance)
+      })
+
+      r3fChildren.push(...threeChildren)
+
+      return r3fChildren
+    }
+
+    return r3fChildren
+  }
 
   public findAll = (decider: (node: ReactThreeTestInstance) => boolean): ReactThreeTestInstance[] =>
-    findAll(this as unknown as ReactThreeTestInstance, decider)
+    findAll(this as unknown as ReactThreeTestInstance, decider, { includeRoot: false })
 
   public find = (decider: (node: ReactThreeTestInstance) => boolean): ReactThreeTestInstance =>
     expectOne(this.findAll(decider), `matching custom checker: ${decider.toString()}`)

packages/test-renderer/src/helpers/graph.ts

@@ -1,5 +1,6 @@
 import type { Instance } from '@react-three/fiber'
 import type { SceneGraphItem } from '../types/public'
+import type * as THREE from 'three'
 
 const graphObjectFactory = (
   type: SceneGraphItem['type'],
@@ -11,5 +12,28 @@ const graphObjectFactory = (
   children,
 })
 
-export const toGraph = (object: Instance): SceneGraphItem[] =>
-  object.children.map((child) => graphObjectFactory(child.object.type, child.object.name ?? '', toGraph(child)))
+// Helper function to process raw THREE.js children objects
+function processThreeChildren(children: THREE.Object3D[]): SceneGraphItem[] {
+  return children.map((object) =>
+    graphObjectFactory(
+      object.type,
+      object.name || '',
+      object.children && object.children.length > 0 ? processThreeChildren(object.children) : [],
+    ),
+  )
+}
+
+export const toGraph = (object: Instance): SceneGraphItem[] => {
+  return object.children.map((child) => {
+    // Process standard R3F children
+    const children = toGraph(child)
+
+    // For primitives, also include THREE.js object children
+    if (child.type === 'primitive' && child.object.children?.length) {
+      const threeChildren = processThreeChildren(child.object.children)
+      children.push(...threeChildren)
+    }
+
+    return graphObjectFactory(child.object.type, child.object.name ?? '', children)
+  })
+}

packages/test-renderer/src/helpers/testInstance.ts

@@ -28,15 +28,31 @@ export const matchProps = (props: Obj, filter: Obj) => {
   return true
 }
 
-export const findAll = (root: ReactThreeTestInstance, decider: (node: ReactThreeTestInstance) => boolean) => {
+interface FindAllOptions {
+  /**
+   * Whether to include the root node in search results.
+   * When false, only searches within children.
+   * @default true
+   */
+  includeRoot?: boolean
+}
+
+export const findAll = (
+  root: ReactThreeTestInstance,
+  decider: (node: ReactThreeTestInstance) => boolean,
+  options: FindAllOptions = { includeRoot: true },
+) => {
   const results = []
 
-  if (decider(root)) {
+  // Only include the root node if the option is enabled
+  if (options.includeRoot !== false && decider(root)) {
     results.push(root)
   }
 
+  // Always search through children
   root.allChildren.forEach((child) => {
-    results.push(...findAll(child, decider))
+    // When recursively searching children, we always want to include their roots
+    results.push(...findAll(child, decider, { includeRoot: true }))
   })
 
   return results