From ca14c5644ca9eb55d7afae9565a0fa8326f44200 Mon Sep 17 00:00:00 2001
From: rmorshea <ryan.morshead@gmail.com>
Date: Wed, 2 Feb 2022 17:41:03 -0800
Subject: [PATCH] mark old state as None if unmounting

---
 src/idom/core/layout.py        | 17 ++++++++++-----
 tests/test_core/test_layout.py | 39 ++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/src/idom/core/layout.py b/src/idom/core/layout.py
index 39c5d55fe..870ede478 100644
--- a/src/idom/core/layout.py
+++ b/src/idom/core/layout.py
@@ -355,11 +355,18 @@ def _render_model_children(
                 else:
                     if old_child_state.is_component_state:
                         self._unmount_model_states([old_child_state])
-                    new_child_state = _update_element_model_state(
-                        old_child_state,
-                        new_state,
-                        index,
-                    )
+                        new_child_state = _make_element_model_state(
+                            new_state,
+                            index,
+                            key,
+                        )
+                        old_child_state = None
+                    else:
+                        new_child_state = _update_element_model_state(
+                            old_child_state,
+                            new_state,
+                            index,
+                        )
                 self._render_model(old_child_state, new_child_state, child)
                 new_children.append(new_child_state.model.current)
                 new_state.children_by_key[key] = new_child_state
diff --git a/tests/test_core/test_layout.py b/tests/test_core/test_layout.py
index c5f0773a6..46fb32112 100644
--- a/tests/test_core/test_layout.py
+++ b/tests/test_core/test_layout.py
@@ -8,6 +8,7 @@
 import pytest
 
 import idom
+from idom import html
 from idom.config import IDOM_DEBUG_MODE
 from idom.core.dispatcher import render_json_patch
 from idom.core.hooks import use_effect
@@ -838,3 +839,41 @@ def HasState():
             root_hook.latest.schedule_render()
             await layout.render()
             assert last_state != state.current
+
+
+async def test_switching_node_type_with_event_handlers():
+    toggle_type = idom.Ref()
+    element_static_handler = StaticEventHandler()
+    component_static_handler = StaticEventHandler()
+
+    @idom.component
+    def Root():
+        toggle, toggle_type.current = use_toggle(True)
+        handler = element_static_handler.use(lambda: None)
+        if toggle:
+            return html.div(html.button({"onEvent": handler}))
+        else:
+            return html.div(SomeComponent())
+
+    @idom.component
+    def SomeComponent():
+        handler = component_static_handler.use(lambda: None)
+        return html.button({"onAnotherEvent": handler})
+
+    with idom.Layout(Root()) as layout:
+        await layout.render()
+
+        assert element_static_handler.target in layout._event_handlers
+        assert component_static_handler.target not in layout._event_handlers
+
+        toggle_type.current()
+        await layout.render()
+
+        assert element_static_handler.target not in layout._event_handlers
+        assert component_static_handler.target in layout._event_handlers
+
+        toggle_type.current()
+        await layout.render()
+
+        assert element_static_handler.target in layout._event_handlers
+        assert component_static_handler.target not in layout._event_handlers