1 files changed, 71 insertions, 0 deletions
diff --git a/app/tests/test_tui.py b/app/tests/test_tui.py
index 49960a1..b206de7 100644
--- a/app/tests/test_tui.py
+++ b/app/tests/test_tui.py
@@ -947,5 +947,76 @@ class FlashTests(TuiTestCase):
         self.assert_inside_border(screen)
 
 
+class WizardTests(unittest.TestCase):
+    """The tui.Wizard screen-stack driver: Esc steps back one screen."""
+
+    def test_advances_to_the_final_value(self):
+        called = []
+
+        def one():
+            called.append("one")
+            return two
+
+        def two():
+            called.append("two")
+            return {"done": True}
+
+        wizard = tui.Wizard()
+        self.assertEqual(wizard.run(one), {"done": True})
+        self.assertEqual(called, ["one", "two"])
+
+    def test_back_on_first_screen_aborts(self):
+        wizard = tui.Wizard()
+
+        def first():
+            return tui.Wizard.BACK
+
+        self.assertIsNone(wizard.run(first))
+
+    def test_back_pops_to_the_previous_screen(self):
+        calls = []
+
+        def first():
+            calls.append("first")
+            # Esc on the re-shown first screen finishes the wizard.
+            return second if calls.count("first") == 1 else {"done": True}
+
+        def second():
+            calls.append("second")
+            return tui.Wizard.BACK
+
+        wizard = tui.Wizard()
+        self.assertEqual(wizard.run(first), {"done": True})
+        self.assertEqual(calls, ["first", "second", "first"])
+
+    def test_back_goes_one_screen_at_a_time(self):
+        order = []
+
+        def a():
+            order.append("a")
+            return b
+
+        def b():
+            order.append("b")
+            return c if order.count("b") == 1 else {"done": True}
+
+        def c():
+            order.append("c")
+            return tui.Wizard.BACK
+
+        wizard = tui.Wizard()
+        self.assertEqual(wizard.run(a), {"done": True})
+        # Back from c lands on b (one screen), not a.
+        self.assertEqual(order, ["a", "b", "c", "b"])
+
+    def test_screen_returning_none_aborts(self):
+        wizard = tui.Wizard()
+
+        def first():
+            return None
+
+        self.assertIsNone(wizard.run(first))
+
+
 if __name__ == "__main__":
     unittest.main()