From 69724315567b4b78681aa009e070844c5ea29981 Mon Sep 17 00:00:00 2001
From: Tomas <tomas.engebretsen@gmail.com>
Date: Thu, 9 Jan 2025 09:16:44 +0100
Subject: [PATCH] feat: Add useUnmount hook

---
 frontend/libs/studio-hooks/src/hooks/index.ts |  1 +
 .../studio-hooks/src/hooks/useUnmount.test.ts | 46 +++++++++++++++++++
 .../libs/studio-hooks/src/hooks/useUnmount.ts | 11 +++++
 3 files changed, 58 insertions(+)
 create mode 100644 frontend/libs/studio-hooks/src/hooks/useUnmount.test.ts
 create mode 100644 frontend/libs/studio-hooks/src/hooks/useUnmount.ts

diff --git a/frontend/libs/studio-hooks/src/hooks/index.ts b/frontend/libs/studio-hooks/src/hooks/index.ts
index 92652568781..1625bae0ed0 100644
--- a/frontend/libs/studio-hooks/src/hooks/index.ts
+++ b/frontend/libs/studio-hooks/src/hooks/index.ts
@@ -2,3 +2,4 @@ export * from './useDebounce';
 export * from './useForwardedRef';
 export * from './usePropState';
 export * from './useUniqueKeys';
+export * from './useUnmount';
diff --git a/frontend/libs/studio-hooks/src/hooks/useUnmount.test.ts b/frontend/libs/studio-hooks/src/hooks/useUnmount.test.ts
new file mode 100644
index 00000000000..f49d66454e0
--- /dev/null
+++ b/frontend/libs/studio-hooks/src/hooks/useUnmount.test.ts
@@ -0,0 +1,46 @@
+import type { RenderHookResult } from '@testing-library/react';
+import { renderHook } from '@testing-library/react';
+import { useUnmount } from './useUnmount';
+
+describe('useUnmount', () => {
+  it('Does not call the function on first render', () => {
+    const fun = jest.fn();
+    renderUseUnmount(fun);
+    expect(fun).not.toHaveBeenCalled();
+  });
+
+  it('Does not call the function on rerender', () => {
+    const fun = jest.fn();
+    const { rerender } = renderUseUnmount(fun);
+    rerender({ fun });
+    expect(fun).not.toHaveBeenCalled();
+  });
+
+  it('Calls the function on unmount', () => {
+    const fun = jest.fn();
+    const { unmount } = renderUseUnmount(fun);
+    unmount();
+    expect(fun).toHaveBeenCalledTimes(1);
+  });
+
+  it('Calls the most recent function on unmount when it has rendered with different functions', () => {
+    const firstFun = jest.fn();
+    const secondFun = jest.fn();
+    const { rerender, unmount } = renderUseUnmount(firstFun);
+    rerender({ fun: secondFun });
+    unmount();
+    expect(firstFun).not.toHaveBeenCalled();
+    expect(secondFun).toHaveBeenCalledTimes(1);
+  });
+});
+
+type RenderUseUnmountProps = {
+  fun: () => void;
+};
+
+function renderUseUnmount(fun: () => void): RenderHookResult<void, RenderUseUnmountProps> {
+  const props: RenderUseUnmountProps = { fun };
+  return renderHook<void, RenderUseUnmountProps>((props) => useUnmount(props.fun), {
+    initialProps: props,
+  });
+}
diff --git a/frontend/libs/studio-hooks/src/hooks/useUnmount.ts b/frontend/libs/studio-hooks/src/hooks/useUnmount.ts
new file mode 100644
index 00000000000..a4025c8937c
--- /dev/null
+++ b/frontend/libs/studio-hooks/src/hooks/useUnmount.ts
@@ -0,0 +1,11 @@
+import { useEffect, useRef } from 'react';
+
+export function useUnmount(fun: () => void): void {
+  const functionRef = useRef<() => void>(fun);
+
+  useEffect(() => {
+    functionRef.current = fun;
+  }, [fun]);
+
+  useEffect(() => () => functionRef.current(), []);
+}