From e4a20927803da34409c5d90fbb8e75e16dce1dca Mon Sep 17 00:00:00 2001
From: Fernando Macedo <fgmacedo@gmail.com>
Date: Wed, 11 Sep 2024 17:33:09 -0300
Subject: [PATCH] fix: Fix #480, engine was not assigned on the initial
 callback (#481)

---
 statemachine/engines/async_.py |  1 +
 statemachine/engines/sync.py   |  1 +
 statemachine/statemachine.py   |  9 +++----
 tests/conftest.py              |  4 ++--
 tests/test_transitions.py      |  4 ++--
 tests/testcases/issue480.md    | 43 ++++++++++++++++++++++++++++++++++
 6 files changed, 54 insertions(+), 8 deletions(-)
 create mode 100644 tests/testcases/issue480.md

diff --git a/statemachine/engines/async_.py b/statemachine/engines/async_.py
index 84310836..b88cce71 100644
--- a/statemachine/engines/async_.py
+++ b/statemachine/engines/async_.py
@@ -15,6 +15,7 @@
 
 class AsyncEngine:
     def __init__(self, sm: "StateMachine", rtc: bool = True):
+        sm._engine = self
         self.sm = proxy(sm)
         self._sentinel = object()
         if not rtc:
diff --git a/statemachine/engines/sync.py b/statemachine/engines/sync.py
index c89c7570..6a658e79 100644
--- a/statemachine/engines/sync.py
+++ b/statemachine/engines/sync.py
@@ -13,6 +13,7 @@
 
 class SyncEngine:
     def __init__(self, sm: "StateMachine", rtc: bool = True):
+        sm._engine = self
         self.sm = proxy(sm)
         self._sentinel = object()
         self._rtc = rtc
diff --git a/statemachine/statemachine.py b/statemachine/statemachine.py
index a2e750aa..fb4e45a1 100644
--- a/statemachine/statemachine.py
+++ b/statemachine/statemachine.py
@@ -105,13 +105,14 @@ def __init__(
             )
             self._put_nonblocking(trigger_data)
 
-        self._engine = self._get_engine(rtc)
+        self._engine: AsyncEngine | SyncEngine | None = None
+        self._select_engine(rtc)
 
-    def _get_engine(self, rtc: bool):
+    def _select_engine(self, rtc: bool):
         if self._callbacks_registry.has_async_callbacks:
-            return AsyncEngine(self, rtc=rtc)
+            AsyncEngine(self, rtc=rtc)
         else:
-            return SyncEngine(self, rtc=rtc)
+            SyncEngine(self, rtc=rtc)
 
     def activate_initial_state(self):
         result = self._engine.activate_initial_state()
diff --git a/tests/conftest.py b/tests/conftest.py
index a720a808..6ad318eb 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -135,8 +135,8 @@ class TrafficLightMachine(StateMachine):
         stop = yellow.to(red)
         go = red.to(green)
 
-        def _get_engine(self, rtc: bool):
-            return engine(self, rtc)
+        def _select_engine(self, rtc: bool):
+            engine(self, rtc)
 
     return TrafficLightMachine
 
diff --git a/tests/test_transitions.py b/tests/test_transitions.py
index c5bd7a81..3b6b52aa 100644
--- a/tests/test_transitions.py
+++ b/tests/test_transitions.py
@@ -252,8 +252,8 @@ class TestStateMachine(StateMachine):
 
             loop = initial.to.itself(internal=internal)
 
-            def _get_engine(self, rtc: bool):
-                return engine(self, rtc)
+            def _select_engine(self, rtc: bool):
+                engine(self, rtc)
 
             def on_exit_initial(self):
                 calls.append("on_exit_initial")
diff --git a/tests/testcases/issue480.md b/tests/testcases/issue480.md
new file mode 100644
index 00000000..b86d37e6
--- /dev/null
+++ b/tests/testcases/issue480.md
@@ -0,0 +1,43 @@
+
+
+### Issue 480
+
+A StateMachine that exercises the example given on issue
+#[480](https://github.com/fgmacedo/python-statemachine/issues/480).
+
+Should be possible to trigger an event on the initial state activation handler.
+
+```py
+>>> from statemachine import StateMachine, State
+>>>
+>>> class MyStateMachine(StateMachine):
+...     State_1 = State(initial=True)
+...     State_2 = State()
+...     Trans_1 = State_1.to(State_2)
+...
+...     def __init__(self):
+...         super(MyStateMachine, self).__init__()
+...
+...     def on_enter_State_1(self):
+...         print("Entering State_1 state")
+...         self.long_running_task()
+...
+...     def on_exit_State_1(self):
+...         print("Exiting State_1 state")
+...
+...     def on_enter_State_2(self):
+...         print("Entering State_2 state")
+...
+...     def long_running_task(self):
+...         print("long running task process started")
+...         self.Trans_1()
+...         print("long running task process ended")
+...
+>>> sm = MyStateMachine()
+Entering State_1 state
+long running task process started
+long running task process ended
+Exiting State_1 state
+Entering State_2 state
+
+```