From 2f03d131ea064e07d65aaa0531d48292b23ae521 Mon Sep 17 00:00:00 2001
From: Piotr Gularski <piotr.gularski@pm.me>
Date: Wed, 17 Jun 2026 13:48:32 +0200
Subject: [PATCH 1/2] Add regression test for issue 8 snapshot restore

---
 test/test_serialization.py | 59 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 59 insertions(+)

diff --git a/test/test_serialization.py b/test/test_serialization.py
index 1cbef9d..f59394c 100644
--- a/test/test_serialization.py
+++ b/test/test_serialization.py
@@ -93,6 +93,65 @@ def test_snapshot_restore_does_not_serialize_user_pda_stack():
     assert list(restored_child.stack.deque) == ['keep child stack']
 
 
+def _issue_8_machine():
+    machine = StateMachine('device')
+    active = StateMachine('active')
+    asleep = State('asleep')
+    idle = State('idle')
+    running = State('running')
+
+    machine.add_state(active, initial=True)
+    machine.add_state(asleep)
+    active.add_state(idle, initial=True)
+    active.add_state(running)
+    active.add_transition(idle, running, events=['start'])
+    machine.add_transition(active, asleep, events=['sleep'])
+    machine.initialize()
+
+    return machine, {
+        'active': active,
+        'asleep': asleep,
+        'idle': idle,
+        'running': running,
+    }
+
+
+def test_issue_8_snapshot_restore_recovers_state_history_after_rebuild():
+    machine, _ = _issue_8_machine()
+    machine.dispatch(Event('start'))
+    machine.dispatch(Event('sleep'))
+
+    data = snapshot(machine)
+
+    restored, restored_states = _issue_8_machine()
+    restore(restored, data)
+
+    assert restored.leaf_state is restored_states['asleep']
+    assert restored.state is restored_states['asleep']
+    assert restored_states['active'].state is restored_states['idle']
+    assert [state.name for state in restored.state_stack.deque] == ['active']
+    assert [state.name for state in
+            restored_states['active'].state_stack.deque] == [
+                'idle',
+                'running',
+            ]
+    assert [state.name for state in restored.leaf_state_stack.deque] == [
+        'idle',
+        'running',
+    ]
+
+    restored.set_previous_leaf_state(Event('wake'))
+
+    assert restored.leaf_state is restored_states['running']
+    assert restored.state is restored_states['active']
+    assert restored_states['active'].state is restored_states['running']
+    assert [state.name for state in restored.leaf_state_stack.deque] == [
+        'idle',
+        'running',
+        'asleep',
+    ]
+
+
 def test_restore_fails_on_topology_mismatch():
     machine = _duplicate_name_machine()
     data = snapshot(machine)

From f9cf24e2974bc70eed895da9750e2a7b5b099f81 Mon Sep 17 00:00:00 2001
From: Piotr Gularski <piotr.gularski@pm.me>
Date: Wed, 17 Jun 2026 13:52:30 +0200
Subject: [PATCH 2/2] Refactor issue 8 serialization test setup

---
 test/test_serialization.py | 49 +++++++++++++++++++-------------------
 1 file changed, 24 insertions(+), 25 deletions(-)

diff --git a/test/test_serialization.py b/test/test_serialization.py
index f59394c..bd1e839 100644
--- a/test/test_serialization.py
+++ b/test/test_serialization.py
@@ -93,37 +93,36 @@ def test_snapshot_restore_does_not_serialize_user_pda_stack():
     assert list(restored_child.stack.deque) == ['keep child stack']
 
 
-def _issue_8_machine():
-    machine = StateMachine('device')
-    active = StateMachine('active')
-    asleep = State('asleep')
-    idle = State('idle')
-    running = State('running')
-
-    machine.add_state(active, initial=True)
-    machine.add_state(asleep)
-    active.add_state(idle, initial=True)
-    active.add_state(running)
-    active.add_transition(idle, running, events=['start'])
-    machine.add_transition(active, asleep, events=['sleep'])
-    machine.initialize()
-
-    return machine, {
-        'active': active,
-        'asleep': asleep,
-        'idle': idle,
-        'running': running,
-    }
-
-
 def test_issue_8_snapshot_restore_recovers_state_history_after_rebuild():
-    machine, _ = _issue_8_machine()
+    def issue_8_machine():
+        machine = StateMachine('device')
+        active = StateMachine('active')
+        asleep = State('asleep')
+        idle = State('idle')
+        running = State('running')
+
+        machine.add_state(active, initial=True)
+        machine.add_state(asleep)
+        active.add_state(idle, initial=True)
+        active.add_state(running)
+        active.add_transition(idle, running, events=['start'])
+        machine.add_transition(active, asleep, events=['sleep'])
+        machine.initialize()
+
+        return machine, {
+            'active': active,
+            'asleep': asleep,
+            'idle': idle,
+            'running': running,
+        }
+
+    machine, _ = issue_8_machine()
     machine.dispatch(Event('start'))
     machine.dispatch(Event('sleep'))
 
     data = snapshot(machine)
 
-    restored, restored_states = _issue_8_machine()
+    restored, restored_states = issue_8_machine()
     restore(restored, data)
 
     assert restored.leaf_state is restored_states['asleep']