From b6bebd39b5e79e23de3a57ff7685a317f557b629 Mon Sep 17 00:00:00 2001
From: Mark Lundin <mlundin@snapchat.com>
Date: Thu, 8 May 2025 15:50:58 +0100
Subject: [PATCH] Add unit tests for Script component and refactor useScript
 hook

- Introduced comprehensive unit tests for the Script component, validating prop passing, ref forwarding, initialization, re-rendering behavior, and cleanup on unmount.
- Refactored the useScript hook to support ref forwarding, enhancing the integration with the Script component.
- Updated type definitions to improve type safety and clarity in the codebase.

These changes enhance the testing coverage and improve the overall functionality of the Script component and its associated hooks.
---
 packages/lib/src/components/Script.test.tsx | 171 ++++++++++++++++++++
 packages/lib/src/components/Script.tsx      |  26 +--
 packages/lib/src/hooks/use-script.tsx       |  29 +++-
 packages/lib/src/utils/types-utils.ts       |   1 +
 4 files changed, 212 insertions(+), 15 deletions(-)
 create mode 100644 packages/lib/src/components/Script.test.tsx

diff --git a/packages/lib/src/components/Script.test.tsx b/packages/lib/src/components/Script.test.tsx
new file mode 100644
index 00000000..4450214c
--- /dev/null
+++ b/packages/lib/src/components/Script.test.tsx
@@ -0,0 +1,171 @@
+import React, { ReactNode, useEffect, useRef } from 'react';
+import { render } from '@testing-library/react';
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { Script } from './Script';
+import { Entity } from '../Entity';
+import { Script as PcScript } from 'playcanvas';
+import { Application } from '../Application';
+import { SubclassOf } from '../utils/types-utils';
+
+const renderWithProviders = (ui: ReactNode) => {
+    return render(
+        <Application>
+            <Entity>
+                {ui}
+            </Entity>
+        </Application>
+    );
+};
+
+describe('Script Component', () => {
+    beforeEach(() => {
+        vi.clearAllMocks();
+        vi.stubEnv('NODE_ENV', 'development');
+    });
+
+    afterEach(() => {
+        vi.unstubAllEnvs();
+    });
+
+    it('should pass props to the script instance', () => {
+        const speed = 2;
+        const str = 'test';
+        const direction = [1, 0, 0];
+
+        class TestingScript extends PcScript {
+            speed: number;
+            direction: number[];
+            str: string;
+
+            initialize() {
+                expect(this.speed).toBe(speed);
+                expect(this.direction).toEqual(direction);
+                expect(this.str).toBe(str);
+            }
+        }
+        
+        renderWithProviders(<Script script={TestingScript} speed={speed} direction={direction} str={str} />);
+    });
+
+    it('should forward ref to script instance', () => {
+
+        class TestScript extends PcScript {}
+
+        const TestComponent = () => {
+            const scriptRef = useRef<SubclassOf<PcScript>>(null);
+
+            useEffect(() => {
+                // Ensure the script instance is created
+                expect(scriptRef.current).toBeInstanceOf(TestScript);
+            }, []);
+            
+            return <Script script={TestScript} ref={scriptRef} />;
+        };
+
+        renderWithProviders(<TestComponent />);
+    });
+
+    describe('Initialization', () => {
+        it('should initialize on first render', () => {
+            const initializeCount = vi.fn();
+
+            class TestScript extends PcScript {
+                initialize() {
+                    initializeCount();
+                }
+            }
+
+            renderWithProviders(<Script script={TestScript} speed={1} />);
+            
+            expect(initializeCount).toHaveBeenCalledTimes(1);
+        });
+    });
+
+    describe('Re-rendering', () => {
+        it('should not re-initialize when props are shallow equal', () => {
+            const initializeCount = vi.fn();
+
+            class TestScript extends PcScript {
+                initialize() {
+                    initializeCount();
+                }
+            }
+
+            const Container = ({ children }: { children: ReactNode }) => (
+                <Application>
+                    <Entity>
+                        {children}
+                    </Entity>
+                </Application>
+            );
+
+            const { rerender } = render(
+                <Container>
+                    <Script script={TestScript} speed={1} />
+                </Container>
+            );
+
+            // Re-render with same props
+            rerender(
+                <Container>
+                    <Script script={TestScript} speed={1} />
+                </Container>
+            );
+
+            expect(initializeCount).toHaveBeenCalledTimes(1);
+        });
+
+        it('should not re-initialize when props change', () => {
+            const initializeCount = vi.fn();
+
+            class TestScript extends PcScript {
+                initialize() {
+                    initializeCount();
+                }
+            }
+
+            const Container = ({ children }: { children: ReactNode }) => (
+                <Application>
+                    <Entity>
+                        {children}
+                    </Entity>
+                </Application>
+            );
+
+            const { rerender } = render(
+                <Container>
+                    <Script script={TestScript} speed={1} />
+                </Container>
+            );
+
+            // Re-render with different props
+            rerender(
+                <Container>
+                    <Script script={TestScript} speed={2} />
+                </Container>
+            );
+
+            expect(initializeCount).toHaveBeenCalledTimes(1);
+        });
+    });
+
+    describe('Cleanup', () => {
+        it('should clean up script instance on unmount', () => {
+            const destroySpy = vi.fn();
+
+            class UnmountScript extends PcScript {
+                initialize() {
+                    this.on('destroy', destroySpy);
+                }
+            }
+
+            const { unmount } = renderWithProviders(
+                <Script script={UnmountScript} />
+            );
+
+            unmount();
+
+            expect(destroySpy).toHaveBeenCalledTimes(1);
+        });
+    });
+}); 
\ No newline at end of file
diff --git a/packages/lib/src/components/Script.tsx b/packages/lib/src/components/Script.tsx
index de200863..0270fbc1 100644
--- a/packages/lib/src/components/Script.tsx
+++ b/packages/lib/src/components/Script.tsx
@@ -1,8 +1,9 @@
-import { AppBase, Entity, Script as PcScript } from "playcanvas";
+import { Script as PcScript } from "playcanvas";
 import { useScript } from "../hooks"
-import { FC, memo, useMemo } from "react";
+import { forwardRef, memo, useMemo } from "react";
 import { ComponentDefinition, validatePropsPartial } from "../utils/validation";
 import { shallowEquals } from "../utils/compare";
+import { SubclassOf } from "../utils/types-utils";
 
 /**
  * The Script component allows you to hook into the entity's lifecycle. This allows you to
@@ -20,18 +21,21 @@ import { shallowEquals } from "../utils/compare";
  * }
  * <Script script={Rotator} />
  */
-const ScriptComponent: FC<ScriptProps> = (props) => {
-
-    const validatedProps = validatePropsPartial(props, componentDefinition, false);
 
+const ScriptComponent = forwardRef<PcScript, ScriptProps>(function ScriptComponent(
+    props,
+    ref
+  ): React.ReactElement | null {
+    const validatedProps = validatePropsPartial(props as ScriptProps, componentDefinition, false);
+  
     const { script, ...restProps } = validatedProps;
-
-    // Memoize props so that the same object reference is passed if props haven't changed
+  
     const memoizedProps = useMemo(() => restProps, [restProps]);
-
-    useScript(script as new (args: { app: AppBase; entity: Entity; }) => PcScript, memoizedProps);
+  
+    useScript(script as SubclassOf<PcScript>, memoizedProps, ref);
+  
     return null;
-};
+  });
 
 // Memoize the component to prevent re-rendering if `script` or `props` are the same
 export const Script = memo(ScriptComponent, (prevProps, nextProps) => {
@@ -39,7 +43,7 @@ export const Script = memo(ScriptComponent, (prevProps, nextProps) => {
 });
 
 interface ScriptProps {
-    script: new (args: { app: AppBase; entity: Entity; }) => PcScript;
+    script: SubclassOf<PcScript>;
     [key: string]: unknown;
 }
 
diff --git a/packages/lib/src/hooks/use-script.tsx b/packages/lib/src/hooks/use-script.tsx
index b6dcf759..037abf82 100644
--- a/packages/lib/src/hooks/use-script.tsx
+++ b/packages/lib/src/hooks/use-script.tsx
@@ -1,7 +1,8 @@
-import { useEffect, useRef } from 'react';
+import { ForwardedRef, useEffect, useRef } from 'react';
 import { useParent } from './use-parent';
 import { useApp } from './use-app';
-import { AppBase, Application, Entity, Script, ScriptComponent } from 'playcanvas';
+import { Application, Entity, Script, ScriptComponent } from 'playcanvas';
+import { SubclassOf } from '../utils/types-utils';
 
 /**
  * This hook is used to create a script component on an entity. 
@@ -13,7 +14,11 @@ import { AppBase, Application, Entity, Script, ScriptComponent } from 'playcanva
  *   myProperty: 'value',
  * });
  */
-export const useScript = (scriptConstructor: new (args: { app: AppBase; entity: Entity; }) => Script, props: Props) : void  => {
+export const useScript = (
+  scriptConstructor: SubclassOf<Script>, 
+  props: Props, 
+  ref: ForwardedRef<Script>
+) : void  => {
   
   const parent: Entity = useParent();
   const app: Application = useApp();
@@ -37,7 +42,15 @@ export const useScript = (scriptConstructor: new (args: { app: AppBase; entity:
       const scriptInstance = scriptComponent.create(scriptConstructor as unknown as typeof Script, {
         properties: { ...props },
         preloading: false,
-      });
+      }) as Script;
+
+      if (ref) {
+        if (typeof ref === 'function') {
+          ref(scriptInstance);
+        } else {
+          ref.current = scriptInstance;
+        }
+      }
 
       scriptRef.current = scriptInstance;
       scriptComponentRef.current = scriptComponent;
@@ -52,6 +65,14 @@ export const useScript = (scriptConstructor: new (args: { app: AppBase; entity:
 
       if (app && app.root && script && scriptComponent) {
         scriptComponent.destroy(scriptName);
+
+        if (ref) {
+          if (typeof ref === 'function') {
+            ref(null);
+          } else {
+            ref.current = null;
+          }
+        }
       }
     };
   }, [app, parent, scriptConstructor]);
diff --git a/packages/lib/src/utils/types-utils.ts b/packages/lib/src/utils/types-utils.ts
index 7f2a69af..9582a96d 100644
--- a/packages/lib/src/utils/types-utils.ts
+++ b/packages/lib/src/utils/types-utils.ts
@@ -33,6 +33,7 @@ export type PublicProps<T> = {
   ]: T[K];
 };
 
+export type SubclassOf<T> = new () => T;
 
 export type Serializable<T> = {
   [K in keyof T]: T[K] extends Color ? string :