Fix hotswap test

Author: Archmonger

tests/test_testing.py

@@ -173,7 +173,7 @@ def test_list_logged_excptions():
         assert logged_errors == [the_error]
 
 
-async def test_hostwap_update_on_change(display: DisplayFixture):
+async def test_hotswap_update_on_change(display: DisplayFixture):
     """Ensure shared hotswapping works
 
     This basically means that previously rendered views of a hotswap component get updated
@@ -183,34 +183,39 @@ async def test_hostwap_update_on_change(display: DisplayFixture):
     hotswap component to be updated
     """
 
-    def make_next_count_constructor(count):
-        """We need to construct a new function so they're different when we set_state"""
+    def hotswap_1():
+        return html.div({"id": "hotswap-1"}, 1)
 
-        def constructor():
-            count.current += 1
-            return html.div({"id": f"hotswap-{count.current}"}, count.current)
+    def hotswap_2():
+        return html.div({"id": "hotswap-2"}, 2)
 
-        return constructor
+    def hotswap_3():
+        return html.div({"id": "hotswap-3"}, 3)
 
     @component
     def ButtonSwapsDivs():
         count = Ref(0)
+        mount, hostswap = _hotswap(update_on_change=True)
 
         async def on_click(event):
-            mount(make_next_count_constructor(count))
-
-        incr = html.button({"on_click": on_click, "id": "incr-button"}, "incr")
-
-        mount, make_hostswap = _hotswap(update_on_change=True)
-        mount(make_next_count_constructor(count))
-        hotswap_view = make_hostswap()
-
-        return html.div(incr, hotswap_view)
+            count.set_current(count.current + 1)
+            if count.current == 1:
+                mount(hotswap_1)
+            if count.current == 2:
+                mount(hotswap_2)
+            if count.current == 3:
+                mount(hotswap_3)
+
+        return html.div(
+            html.button({"on_click": on_click, "id": "incr-button"}, "incr"),
+            hostswap(),
+        )
 
     await display.show(ButtonSwapsDivs)
 
     client_incr_button = await display.page.wait_for_selector("#incr-button")
 
+    await client_incr_button.click()
     await display.page.wait_for_selector("#hotswap-1")
     await client_incr_button.click()
     await display.page.wait_for_selector("#hotswap-2")