refactor: use dataclass for StateMachineEvent and simplify mock classes

agentspype/agent/publishing.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 from enum import Enum
 from typing import TYPE_CHECKING, Any
 
@@ -27,6 +28,7 @@ class StateAgentPublishing(AgentPublishing):
     class Events(Enum):
         StateMachineTransition = "sm_transition_event"
 
+    @dataclass
     class StateMachineEvent:
         event: Any
         new_state: State

tests/test_agency.py

@@ -58,20 +58,7 @@ def __init__(self, agent: Agent) -> None:
     def trigger_event(self, event_publication: Any, event_data: Any) -> None:
         self.published_events.append((event_publication, event_data))
 
-    class StateMachineEvent:
-        def __init__(self, event: str, state: State) -> None:
-            self.event = event
-            self.new_state = state
-
-    # Override the event publication to use our mock event class
-    sm_transition_event = type(
-        "EventPublication",
-        (),
-        {
-            "event_tag": StateAgentPublishing.Events.StateMachineTransition,
-            "event_class": StateMachineEvent,
-        },
-    )()
+    sm_transition_event = StateAgentPublishing.sm_transition_event
 
 
 class MockAgent(Agent):

tests/test_agent.py

@@ -89,20 +89,8 @@ def __init__(self, agent: Agent) -> None:
     def trigger_event(self, event_publication: Any, event_data: Any) -> None:
         self.published_events.append((event_publication, event_data))
 
-    class StateMachineEvent:
-        def __init__(self, event: str, state: State) -> None:
-            self.event = event
-            self.new_state = state
-
     # Override the event publication to use our mock event class
-    sm_transition_event = type(
-        "EventPublication",
-        (),
-        {
-            "event_tag": StateAgentPublishing.Events.StateMachineTransition,
-            "event_class": StateMachineEvent,
-        },
-    )()
+    sm_transition_event = StateAgentPublishing.sm_transition_event
 
 
 # Define agents with different state machines