Handle onmount callbacks correctly when render conditionally

integration-tests/hooks.test.ts

@@ -1,4 +1,4 @@
-import { screen, waitFor } from "@testing-library/dom";
+import { screen } from "@testing-library/dom";
 import userEvent from "@testing-library/user-event";
 
 import {
@@ -7,6 +7,7 @@ import {
   createState,
   onUnmount,
   onMount,
+  type State,
 } from "../src";
 
 describe("lifecycle hooks", () => {
@@ -147,4 +148,173 @@ describe("lifecycle hooks", () => {
 
     expect(isElementFound).toBe(true);
   });
+
+  test("onMount hooks are executed when new components are mounted based on state", async () => {
+    function App() {
+      const showState = createState(false);
+      return createElement("div", {
+        "data-testid": "appComponent",
+        children: [
+          "app component",
+          createElement("button", {
+            "data-testid": "button",
+            onClick: () => showState.setValue(() => true),
+          }),
+          showState.useValue((shouldShow) =>
+            shouldShow ? createElement(Wrapper) : null
+          ),
+        ],
+      });
+    }
+
+    function Wrapper() {
+      return createElement("div", {
+        children: ["yo", createElement(ConditionalComponent)],
+      });
+    }
+
+    const spy = jest.fn();
+    function ConditionalComponent() {
+      onMount(spy);
+
+      return createElement("div", {
+        children: "conditional component",
+      });
+    }
+
+    cleanup = attachComponent({
+      htmlElement: document.body,
+      component: createElement(App),
+    });
+
+    // hacky way to wait until the next tick so that mount hooks are executed
+    await new Promise((resolve) => {
+      setTimeout(resolve, 0);
+    });
+
+    expect(spy).not.toHaveBeenCalled();
+
+    await userEvent.click(screen.getByTestId("button"));
+
+    // hacky way to wait until the next tick so that mount hooks are executed
+    await new Promise((resolve) => {
+      setTimeout(resolve, 0);
+    });
+
+    expect(spy).toHaveBeenCalledTimes(1);
+  });
+
+  test("onMount is called one time", async () => {
+    const appMountSpy = jest.fn();
+    function App() {
+      onMount(appMountSpy);
+      return createElement("div", {
+        "data-testid": "appComponent",
+        children: ["app component", createElement(FirstComponent)],
+      });
+    }
+
+    const firstComponentMountSpy = jest.fn();
+    function FirstComponent() {
+      onMount(firstComponentMountSpy);
+
+      return createElement("div", {
+        children: ["test", createElement(SecondComponent)],
+      });
+    }
+
+    const secondComponentMountSpy = jest.fn();
+    function SecondComponent() {
+      onMount(secondComponentMountSpy);
+
+      return createElement(ThirdComponent);
+    }
+
+    const thirdComponentMountSpy = jest.fn();
+    function ThirdComponent() {
+      onMount(thirdComponentMountSpy);
+
+      return createElement("div", { children: "test" });
+    }
+
+    cleanup = attachComponent({
+      htmlElement: document.body,
+      component: createElement(App),
+    });
+
+    // hacky way to wait until the next tick so that mount hooks are executed
+    await new Promise((resolve) => {
+      setTimeout(resolve, 0);
+    });
+
+    expect(appMountSpy).toHaveBeenCalledTimes(1);
+    expect(firstComponentMountSpy).toHaveBeenCalledTimes(1);
+    expect(secondComponentMountSpy).toHaveBeenCalledTimes(1);
+    expect(thirdComponentMountSpy).toHaveBeenCalledTimes(1);
+  });
+
+  test("onMount is called correctly for iterator values", async () => {
+    function App() {
+      const tasksState = createState([
+        { id: 1, title: "first task" },
+        { id: 2, title: "second task" },
+      ]);
+      return createElement("div", {
+        "data-testid": "appComponent",
+        children: [
+          "app component",
+          createElement("button", {
+            "data-testid": "button",
+            onClick: () => {
+              tasksState.setValue((currentTasks) =>
+                currentTasks.concat({ id: 3, title: "third task" })
+              );
+            },
+          }),
+          tasksState.useValueIterator({ key: "id" }, ({ elementState }) =>
+            createElement(Task, { taskState: elementState })
+          ),
+        ],
+      });
+    }
+
+    const taskMountSpy = jest.fn();
+    function Task({
+      taskState,
+    }: {
+      taskState: State<{ id: number; title: string }>;
+    }) {
+      onMount(taskMountSpy);
+      return createElement("div", {
+        children: [
+          "task",
+          taskState.useValueSelector(
+            (task) => task.title,
+            (title) => title
+          ),
+        ],
+      });
+    }
+
+    cleanup = attachComponent({
+      htmlElement: document.body,
+      component: createElement(App),
+    });
+
+    // hacky way to wait until the next tick so that mount hooks are executed
+    await new Promise((resolve) => {
+      setTimeout(resolve, 0);
+    });
+
+    expect(taskMountSpy).toHaveBeenCalledTimes(2);
+
+    await userEvent.click(screen.getByTestId("button"));
+
+    // hacky way to wait until the next tick so that mount hooks are executed
+    await new Promise((resolve) => {
+      setTimeout(resolve, 0);
+    });
+
+    expect(taskMountSpy).toHaveBeenCalledTimes(3);
+  });
 });

src/attach-component.ts

@@ -1,4 +1,4 @@
-import { getComponentVelesNode } from "./utils";
+import { getComponentVelesNode, callMountHandlers } from "./utils";
 import { createElement } from "./create-element";
 
 import type { VelesElement, VelesComponent } from "./types";
@@ -16,6 +16,7 @@ function attachComponent({
   const wrappedApp = createElement("div", { children: [component] });
   const { velesElementNode } = getComponentVelesNode(wrappedApp);
   htmlElement.appendChild(velesElementNode.html);
+  callMountHandlers(wrappedApp);
 
   // TODO: iterate over every child and call their `onUnmout` method
   // and add tests for that

src/create-element/parse-children.ts

@@ -44,19 +44,15 @@ function parseChildren({
               htmlElement.append(childComponentofPhantom.html);
               childComponentofPhantom.parentVelesElement = velesNode;
             } else {
-              const { componentsTree, velesElementNode } =
-                getComponentVelesNode(childComponentofPhantom);
+              const { velesElementNode } = getComponentVelesNode(
+                childComponentofPhantom
+              );
 
               if (!velesElementNode) {
                 console.error("can't find HTML tree in a component chain");
               } else {
                 htmlElement.append(velesElementNode.html);
 
-                // TODO: address the same concern as below
-                componentsTree.forEach((component) => {
-                  component._privateMethods._callMountHandlers();
-                });
-
                 velesElementNode.parentVelesElement = velesNode;
               }
             }
@@ -98,13 +94,6 @@ function parseChildren({
                     console.error("can't find HTML tree in a component chain");
                   } else {
                     htmlElement.append(velesElementNode.html);
-
-                    // Same explanation as below. Components are mounted synchronously
-                    setTimeout(() => {
-                      componentsTree.forEach((component) => {
-                        component._privateMethods._callMountHandlers();
-                      });
-                    }, 0);
                     velesElementNode.parentVelesElement = velesNode;
                   }
                 }
@@ -114,15 +103,6 @@ function parseChildren({
             htmlElement.append(velesElementNode.html);
           }
 
-          /**
-           * Components are mounted synchronously, so we can safely wait for the next
-           * CPU tick and be sure that new markup is attached to DOM.
-           */
-          setTimeout(() => {
-            componentsTree.forEach((component) => {
-              component._privateMethods._callMountHandlers();
-            });
-          }, 0);
           velesElementNode.parentVelesElement = velesNode;
           childComponents.push(childComponent);
         }

src/hooks/create-state.ts

@@ -1,4 +1,4 @@
-import { getComponentVelesNode, identity } from "../utils";
+import { getComponentVelesNode, callMountHandlers, identity } from "../utils";
 import { onUnmount, onMount } from "./lifecycle";
 import { createElement } from "../create-element/create-element";
 import { createTextElement } from "../create-element/create-text-element";
@@ -380,6 +380,9 @@ function createState<T>(
               );
             // we call unmount handlers right after we replace it
             node._privateMethods._callUnmountHandlers();
+            // at this point the new Node is mounted, childComponents are updated
+            // and unmount handlers for the old node are called
+            callMountHandlers(newNode);
 
             // right after that, we add the callback back
             // the top level node is guaranteed to be rendered again (at least right now)
@@ -648,6 +651,8 @@ function createState<T>(
               offset = offset + 1;
               currentElement = newNodeVelesElement.html;
               newElementsCount = newElementsCount + 1;
+
+              callMountHandlers(newNode);
             }
           });
 

src/utils.ts

@@ -33,8 +33,28 @@ function getComponentVelesNode(
   return { velesElementNode: childNode, componentsTree };
 }
 
+function callMountHandlers(
+  component: VelesComponent | VelesElement | VelesStringElement
+): void {
+  if ("velesStringElement" in component) {
+    // string elements don't have mount callbacks, only unmount
+    return;
+  }
+
+  if ("velesComponent" in component) {
+    component._privateMethods._callMountHandlers();
+    callMountHandlers(component.tree);
+  }
+
+  if ("velesNode" in component) {
+    component.childComponents.forEach((childComponent) =>
+      callMountHandlers(childComponent)
+    );
+  }
+}
+
 function identity<T>(value1: T, value2: T) {
   return value1 === value2;
 }
 
-export { getComponentVelesNode, identity };
+export { getComponentVelesNode, identity, callMountHandlers };