feat(events): improve event mapping

hathor/cli/events_simulator/scenario.py

@@ -16,7 +16,7 @@
 
 from hathor.event.model.base_event import BaseEvent
 from hathor.event.model.event_data import TxData, TxMetadata
-from hathor.pubsub import HathorEvents
+from hathor.event.model.event_type import EventType
 
 
 class Scenario(Enum):
@@ -57,7 +57,7 @@ def get_events(self):
     peer_id='123',
     id=0,
     timestamp=0,
-    type=HathorEvents.NETWORK_NEW_TX_ACCEPTED,
+    type=EventType.NEW_VERTEX_ACCEPTED,
     data=_TRANSACTION_DATA_1
 )
 

hathor/event/event_manager.py

@@ -17,6 +17,7 @@
 from structlog import get_logger
 
 from hathor.event.model.base_event import BaseEvent
+from hathor.event.model.event_type import EventType
 from hathor.event.storage import EventStorage
 from hathor.event.websocket import EventWebsocketFactory
 from hathor.pubsub import EventArguments, HathorEvents, PubSubManager
@@ -25,27 +26,22 @@
 logger = get_logger()
 
 _GROUP_START_EVENTS = {
-    HathorEvents.REORG_STARTED,
+    EventType.REORG_STARTED,
 }
 
 _GROUP_END_EVENTS = {
-    HathorEvents.REORG_FINISHED,
+    EventType.REORG_FINISHED,
 }
 
 _SUBSCRIBE_EVENTS = [
-    HathorEvents.NETWORK_NEW_TX_ACCEPTED,
+    HathorEvents.MANAGER_ON_START,
     HathorEvents.LOAD_FINISHED,
+    HathorEvents.NETWORK_NEW_TX_ACCEPTED,
     HathorEvents.REORG_STARTED,
     HathorEvents.REORG_FINISHED,
-    HathorEvents.VERTEX_METADATA_CHANGED,
     HathorEvents.CONSENSUS_TX_UPDATE,
-    HathorEvents.CONSENSUS_TX_REMOVED,
 ]
 
-_EVENT_CONVERTERS = {
-    HathorEvents.CONSENSUS_TX_UPDATE: HathorEvents.VERTEX_METADATA_CHANGED
-}
-
 
 class EventManager:
     """Class that manages integration events.
@@ -107,7 +103,7 @@ def _event_group_is_closed(self):
         return (
             self._last_event is None or
             self._last_event.group_id is None or
-            HathorEvents(self._last_event.type) in _GROUP_END_EVENTS
+            EventType(self._last_event.type) in _GROUP_END_EVENTS
         )
 
     def _subscribe_events(self):
@@ -116,12 +112,12 @@ def _subscribe_events(self):
         for event in _SUBSCRIBE_EVENTS:
             self._pubsub.subscribe(event, self._handle_event)
 
-    def _handle_event(self, event_type: HathorEvents, event_args: EventArguments) -> None:
+    def _handle_event(self, hathor_event: HathorEvents, event_args: EventArguments) -> None:
         assert self._is_running, 'Cannot handle event, EventManager is not started.'
 
-        event_type = _EVENT_CONVERTERS.get(event_type, event_type)
+        event_type = EventType.from_hathor_event(hathor_event)
         event_specific_handlers = {
-            HathorEvents.LOAD_FINISHED: self._handle_load_finished
+            EventType.LOAD_FINISHED: self._handle_load_finished
         }
 
         if event_specific_handler := event_specific_handlers.get(event_type):
@@ -132,8 +128,8 @@ def _handle_event(self, event_type: HathorEvents, event_args: EventArguments) ->
 
         self._handle_event_creation(event_type, event_args)
 
-    def _handle_event_creation(self, event_type: HathorEvents, event_args: EventArguments) -> None:
-        create_event_fn: Callable[[HathorEvents, EventArguments], BaseEvent]
+    def _handle_event_creation(self, event_type: EventType, event_args: EventArguments) -> None:
+        create_event_fn: Callable[[EventType, EventArguments], BaseEvent]
 
         if event_type in _GROUP_START_EVENTS:
             create_event_fn = self._create_group_start_event
@@ -149,7 +145,7 @@ def _handle_event_creation(self, event_type: HathorEvents, event_args: EventArgu
 
         self._last_event = event
 
-    def _create_group_start_event(self, event_type: HathorEvents, event_args: EventArguments) -> BaseEvent:
+    def _create_group_start_event(self, event_type: EventType, event_args: EventArguments) -> BaseEvent:
         assert self._event_group_is_closed(), 'A new event group cannot be started as one is already in progress.'
 
         new_group_id = 0 if self._last_existing_group_id is None else self._last_existing_group_id + 1
@@ -162,7 +158,7 @@ def _create_group_start_event(self, event_type: HathorEvents, event_args: EventA
             group_id=new_group_id,
         )
 
-    def _create_group_end_event(self, event_type: HathorEvents, event_args: EventArguments) -> BaseEvent:
+    def _create_group_end_event(self, event_type: EventType, event_args: EventArguments) -> BaseEvent:
         assert self._last_event is not None, 'Cannot end event group if there are no events.'
         assert not self._event_group_is_closed(), 'Cannot end event group as none is in progress.'
 
@@ -172,7 +168,7 @@ def _create_group_end_event(self, event_type: HathorEvents, event_args: EventArg
             group_id=self._last_event.group_id,
         )
 
-    def _create_non_group_edge_event(self, event_type: HathorEvents, event_args: EventArguments) -> BaseEvent:
+    def _create_non_group_edge_event(self, event_type: EventType, event_args: EventArguments) -> BaseEvent:
         group_id = None
 
         if not self._event_group_is_closed():
@@ -190,7 +186,7 @@ def _handle_load_finished(self):
 
     def _create_event(
         self,
-        event_type: HathorEvents,
+        event_type: EventType,
         event_args: EventArguments,
         group_id: Optional[int],
     ) -> BaseEvent:

hathor/event/model/base_event.py

@@ -12,24 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, Optional, Type
+from typing import Optional
 
 from pydantic import NonNegativeInt, validator
 
-from hathor.event.model.event_data import BaseEventData, EmptyData, EventData, ReorgData, TxData
-from hathor.pubsub import EventArguments, HathorEvents
+from hathor.event.model.event_data import EventData
+from hathor.event.model.event_type import EventType
+from hathor.pubsub import EventArguments
 from hathor.utils.pydantic import BaseModel
 
-_EVENT_DATA_MAP: Dict[HathorEvents, Type[BaseEventData]] = {
-    HathorEvents.LOAD_FINISHED: EmptyData,
-    HathorEvents.NETWORK_NEW_TX_ACCEPTED: TxData,
-    HathorEvents.REORG_STARTED: ReorgData,
-    HathorEvents.REORG_FINISHED: EmptyData,
-    HathorEvents.VERTEX_METADATA_CHANGED: TxData,
-    HathorEvents.CONSENSUS_TX_UPDATE: TxData,
-    HathorEvents.CONSENSUS_TX_REMOVED: TxData,
-}
-
 
 class BaseEvent(BaseModel, use_enum_values=True):
     # Full node id, because different full nodes can have different sequences of events
@@ -41,7 +32,7 @@ class BaseEvent(BaseModel, use_enum_values=True):
     # events it's possible that timestamps will temporarily decrease.
     timestamp: float
     # One of the event types
-    type: HathorEvents
+    type: EventType
     # Variable for event type
     data: EventData
     # Used to link events, for example, many TX_METADATA_CHANGED will have the same group_id when they belong to the
@@ -54,14 +45,11 @@ def from_event_arguments(
         peer_id: str,
         event_id: NonNegativeInt,
         timestamp: float,
-        event_type: HathorEvents,
+        event_type: EventType,
         event_args: EventArguments,
         group_id: Optional[NonNegativeInt]
     ) -> 'BaseEvent':
-        event_data_type = _EVENT_DATA_MAP.get(event_type)
-
-        if event_data_type is None:
-            raise ValueError(f'The given event type ({event_type}) is not a supported event')
+        event_data_type = event_type.data_type()
 
         return cls(
             peer_id=peer_id,
@@ -74,8 +62,8 @@ def from_event_arguments(
 
     @validator('data')
     def data_type_must_match_event_type(cls, v, values):
-        event_type = HathorEvents(values['type'])
-        expected_data_type = _EVENT_DATA_MAP.get(event_type)
+        event_type = EventType(values['type'])
+        expected_data_type = event_type.data_type()
 
         if type(v) != expected_data_type:
             raise ValueError('event data type does not match event type')

hathor/event/model/event_type.py

@@ -0,0 +1,58 @@
+#  Copyright 2023 Hathor Labs
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from enum import Enum
+from typing import Dict, Type
+
+from hathor.event.model.event_data import BaseEventData, EmptyData, ReorgData, TxData
+from hathor.pubsub import HathorEvents
+
+
+class EventType(Enum):
+    LOAD_STARTED = 'LOAD_STARTED'
+    LOAD_FINISHED = 'LOAD_FINISHED'
+    NEW_VERTEX_ACCEPTED = 'NEW_VERTEX_ACCEPTED'
+    REORG_STARTED = 'REORG_STARTED'
+    REORG_FINISHED = 'REORG_FINISHED'
+    VERTEX_METADATA_CHANGED = 'VERTEX_METADATA_CHANGED'
+
+    @classmethod
+    def from_hathor_event(cls, hathor_event: HathorEvents) -> 'EventType':
+        event = _HATHOR_EVENT_TO_EVENT_TYPE.get(hathor_event)
+
+        assert event is not None, f'Cannot create EventType from {hathor_event}'
+
+        return event
+
+    def data_type(self) -> Type[BaseEventData]:
+        return _EVENT_TYPE_TO_EVENT_DATA[self]
+
+
+_HATHOR_EVENT_TO_EVENT_TYPE = {
+    HathorEvents.MANAGER_ON_START: EventType.LOAD_STARTED,
+    HathorEvents.LOAD_FINISHED: EventType.LOAD_FINISHED,
+    HathorEvents.NETWORK_NEW_TX_ACCEPTED: EventType.NEW_VERTEX_ACCEPTED,
+    HathorEvents.REORG_STARTED: EventType.REORG_STARTED,
+    HathorEvents.REORG_FINISHED: EventType.REORG_FINISHED,
+    HathorEvents.CONSENSUS_TX_UPDATE: EventType.VERTEX_METADATA_CHANGED
+}
+
+_EVENT_TYPE_TO_EVENT_DATA: Dict[EventType, Type[BaseEventData]] = {
+    EventType.LOAD_STARTED: EmptyData,
+    EventType.LOAD_FINISHED: EmptyData,
+    EventType.NEW_VERTEX_ACCEPTED: TxData,
+    EventType.REORG_STARTED: ReorgData,
+    EventType.REORG_FINISHED: EmptyData,
+    EventType.VERTEX_METADATA_CHANGED: TxData,
+}

hathor/pubsub.py

@@ -135,8 +135,6 @@ class HathorEvents(Enum):
 
     REORG_FINISHED = 'reorg:finished'
 
-    VERTEX_METADATA_CHANGED = 'vertex:metadata_changed'
-
 
 class EventArguments:
     """Simple object for storing event arguments.

tests/event/test_base_event.py

@@ -17,7 +17,7 @@
 
 from hathor.event.model.base_event import BaseEvent
 from hathor.event.model.event_data import ReorgData
-from hathor.pubsub import HathorEvents
+from hathor.event.model.event_type import EventType
 from tests.utils import EventMocker
 
 
@@ -28,7 +28,7 @@ def test_create_base_event(event_id, group_id):
         peer_id='some_peer',
         id=event_id,
         timestamp=123.3,
-        type=HathorEvents.VERTEX_METADATA_CHANGED,
+        type=EventType.VERTEX_METADATA_CHANGED,
         data=EventMocker.tx_data,
         group_id=group_id
     )
@@ -37,7 +37,7 @@ def test_create_base_event(event_id, group_id):
         peer_id='some_peer',
         id=event_id,
         timestamp=123.3,
-        type='vertex:metadata_changed',
+        type='VERTEX_METADATA_CHANGED',
         data=dict(
             hash='abc',
             nonce=123,
@@ -78,7 +78,7 @@ def test_create_base_event_fail_id(event_id):
             peer_id='some_peer',
             id=event_id,
             timestamp=123.3,
-            type=HathorEvents.VERTEX_METADATA_CHANGED,
+            type=EventType.VERTEX_METADATA_CHANGED,
             data=EventMocker.tx_data,
         )
 
@@ -90,7 +90,7 @@ def test_create_base_event_fail_group_id(group_id):
             peer_id='some_peer',
             id=0,
             timestamp=123.3,
-            type=HathorEvents.VERTEX_METADATA_CHANGED,
+            type=EventType.VERTEX_METADATA_CHANGED,
             data=EventMocker.tx_data,
             group_id=group_id
         )
@@ -102,7 +102,7 @@ def test_create_base_event_fail_data_type():
             peer_id='some_peer',
             id=0,
             timestamp=123.3,
-            type=HathorEvents.VERTEX_METADATA_CHANGED,
+            type=EventType.VERTEX_METADATA_CHANGED,
             data=ReorgData(
                 reorg_size=10,
                 previous_best_block='a',

tests/event/test_event_manager.py

@@ -1,6 +1,7 @@
 from unittest.mock import Mock
 
 from hathor.event import EventManager
+from hathor.event.model.event_type import EventType
 from hathor.event.storage.memory_storage import EventMemoryStorage
 from hathor.event.websocket import EventWebsocketFactory
 from hathor.pubsub import HathorEvents, PubSubManager
@@ -55,10 +56,10 @@ def test_event_group(self):
         event2 = self.event_storage.get_event(2)
         event3 = self.event_storage.get_event(3)
         event4 = self.event_storage.get_event(4)
-        self.assertEqual(HathorEvents(event0.type), HathorEvents.LOAD_FINISHED)
-        self.assertEqual(HathorEvents(event1.type), HathorEvents.REORG_STARTED)
+        self.assertEqual(EventType(event0.type), EventType.LOAD_FINISHED)
+        self.assertEqual(EventType(event1.type), EventType.REORG_STARTED)
         self.assertIsNotNone(event1.group_id)
-        self.assertEqual(HathorEvents(event2.type), HathorEvents.REORG_FINISHED)
+        self.assertEqual(EventType(event2.type), EventType.REORG_FINISHED)
         self.assertIsNotNone(event2.group_id)
         self.assertEqual(event1.group_id, event2.group_id)
         self.assertNotEqual(event2.group_id, event3.group_id)