From 6556fb73c061a71ec88a477046dfb7dd6d72ec6c Mon Sep 17 00:00:00 2001
From: Seva Zaikov <seva.zaikov@gmail.com>
Date: Sun, 9 Jun 2024 00:32:57 -0700
Subject: [PATCH] call unmount callbacks on entire Veles app removal (#53)

---
 integration-tests/attach-component.test.ts | 55 +++++++++++++++++++++-
 src/attach-component.ts                    |  1 +
 src/create-element/create-element.ts       |  7 +--
 3 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/integration-tests/attach-component.test.ts b/integration-tests/attach-component.test.ts
index 6bd9f63..c449b12 100644
--- a/integration-tests/attach-component.test.ts
+++ b/integration-tests/attach-component.test.ts
@@ -1,10 +1,16 @@
 import { screen } from "@testing-library/dom";
 
-import { attachComponent, createElement } from "../src";
+import {
+  attachComponent,
+  createElement,
+  createState,
+  onUnmount,
+  Fragment,
+} from "../src";
 
 describe("attachComponent", () => {
   test("attaches component tree correctly", () => {
-    attachComponent({
+    const removeVelesTree = attachComponent({
       htmlElement: document.body,
       component: createElement("div", {
         "data-testid": "attachedComponent",
@@ -14,5 +20,50 @@ describe("attachComponent", () => {
 
     const docElement = screen.getByTestId("attachedComponent");
     expect(docElement).toBeVisible();
+
+    removeVelesTree();
+  });
+
+  test("calls all onUnmount callbacks when removing tree", () => {
+    const appUnmountSpy = jest.fn();
+    function App() {
+      onUnmount(appUnmountSpy);
+
+      return createElement("div", {
+        children: [
+          createElement(Fragment, {
+            children: [createElement(Child), createElement(Child)],
+          }),
+        ],
+      });
+    }
+    const state = createState(0);
+    const childUnmountSpy = jest.fn();
+    const childSubscriptionSpy = jest.fn();
+    function Child() {
+      onUnmount(childUnmountSpy);
+      state.trackValue(childSubscriptionSpy, { skipFirstCall: true });
+      return createElement("div", {
+        children: "test",
+      });
+    }
+    const removeVelesTree = attachComponent({
+      htmlElement: document.body,
+      component: createElement("div", {
+        "data-testid": "attachedComponent",
+        children: [createElement(App)],
+      }),
+    });
+
+    state.setValue(1);
+    expect(childSubscriptionSpy).toHaveBeenCalledTimes(2);
+
+    removeVelesTree();
+    expect(appUnmountSpy).toHaveBeenCalledTimes(1);
+    expect(childUnmountSpy).toHaveBeenCalledTimes(2);
+
+    state.setValue(2);
+    // subscriptions were removed, so no more calls
+    expect(childSubscriptionSpy).toHaveBeenCalledTimes(2);
   });
 });
diff --git a/src/attach-component.ts b/src/attach-component.ts
index 2847f3a..3a34409 100644
--- a/src/attach-component.ts
+++ b/src/attach-component.ts
@@ -21,6 +21,7 @@ function attachComponent({
   // TODO: iterate over every child and call their `onUnmout` method
   // and add tests for that
   return () => {
+    wrappedApp._privateMethods._callUnmountHandlers();
     velesElementNode.html.remove();
   };
 }
diff --git a/src/create-element/create-element.ts b/src/create-element/create-element.ts
index 90d2821..5eb3aed 100644
--- a/src/create-element/create-element.ts
+++ b/src/create-element/create-element.ts
@@ -34,12 +34,13 @@ function createElement(
     // `useAttribute`
     let unmountHandlers: Function[] = [];
     const callUnmountHandlers = () => {
-      unmountHandlers.forEach((cb) => cb());
-      unmountHandlers = [];
-
+      // `onUnmount` is logically better to be executed on children first
       childComponents.forEach((childComponent) => {
         childComponent._privateMethods._callUnmountHandlers();
       });
+
+      unmountHandlers.forEach((cb) => cb());
+      unmountHandlers = [];
     };
 
     velesNode.html = newElement;