From 852bb1922348b979fca6195d6fb1f82cb1058fc3 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sat, 29 Jun 2024 07:49:16 -0500
Subject: [PATCH] Cleanup db_schema from_event constructors (#120803)

---
 .../components/recorder/db_schema.py          | 73 ++++++++++---------
 homeassistant/core.py                         |  5 ++
 tests/test_core.py                            |  8 ++
 3 files changed, 50 insertions(+), 36 deletions(-)

diff --git a/homeassistant/components/recorder/db_schema.py b/homeassistant/components/recorder/db_schema.py
index ba4a6106bce..915fd4a4bb8 100644
--- a/homeassistant/components/recorder/db_schema.py
+++ b/homeassistant/components/recorder/db_schema.py
@@ -238,7 +238,6 @@ class JSONLiteral(JSON):
 
 
 EVENT_ORIGIN_ORDER = [EventOrigin.local, EventOrigin.remote]
-EVENT_ORIGIN_TO_IDX = {origin: idx for idx, origin in enumerate(EVENT_ORIGIN_ORDER)}
 
 
 class Events(Base):
@@ -305,18 +304,19 @@ class Events(Base):
     @staticmethod
     def from_event(event: Event) -> Events:
         """Create an event database object from a native event."""
+        context = event.context
         return Events(
             event_type=None,
             event_data=None,
-            origin_idx=EVENT_ORIGIN_TO_IDX.get(event.origin),
+            origin_idx=event.origin.idx,
             time_fired=None,
             time_fired_ts=event.time_fired_timestamp,
             context_id=None,
-            context_id_bin=ulid_to_bytes_or_none(event.context.id),
+            context_id_bin=ulid_to_bytes_or_none(context.id),
             context_user_id=None,
-            context_user_id_bin=uuid_hex_to_bytes_or_none(event.context.user_id),
+            context_user_id_bin=uuid_hex_to_bytes_or_none(context.user_id),
             context_parent_id=None,
-            context_parent_id_bin=ulid_to_bytes_or_none(event.context.parent_id),
+            context_parent_id_bin=ulid_to_bytes_or_none(context.parent_id),
         )
 
     def to_native(self, validate_entity_id: bool = True) -> Event | None:
@@ -492,41 +492,42 @@ class States(Base):
     @staticmethod
     def from_event(event: Event[EventStateChangedData]) -> States:
         """Create object from a state_changed event."""
-        entity_id = event.data["entity_id"]
         state = event.data["new_state"]
-        dbstate = States(
-            entity_id=entity_id,
-            attributes=None,
-            context_id=None,
-            context_id_bin=ulid_to_bytes_or_none(event.context.id),
-            context_user_id=None,
-            context_user_id_bin=uuid_hex_to_bytes_or_none(event.context.user_id),
-            context_parent_id=None,
-            context_parent_id_bin=ulid_to_bytes_or_none(event.context.parent_id),
-            origin_idx=EVENT_ORIGIN_TO_IDX.get(event.origin),
-            last_updated=None,
-            last_changed=None,
-        )
         # None state means the state was removed from the state machine
         if state is None:
-            dbstate.state = ""
-            dbstate.last_updated_ts = event.time_fired_timestamp
-            dbstate.last_changed_ts = None
-            dbstate.last_reported_ts = None
-            return dbstate
-
-        dbstate.state = state.state
-        dbstate.last_updated_ts = state.last_updated_timestamp
-        if state.last_updated == state.last_changed:
-            dbstate.last_changed_ts = None
+            state_value = ""
+            last_updated_ts = event.time_fired_timestamp
+            last_changed_ts = None
+            last_reported_ts = None
         else:
-            dbstate.last_changed_ts = state.last_changed_timestamp
-        if state.last_updated == state.last_reported:
-            dbstate.last_reported_ts = None
-        else:
-            dbstate.last_reported_ts = state.last_reported_timestamp
-
-        return dbstate
+            state_value = state.state
+            last_updated_ts = state.last_updated_timestamp
+            if state.last_updated == state.last_changed:
+                last_changed_ts = None
+            else:
+                last_changed_ts = state.last_changed_timestamp
+            if state.last_updated == state.last_reported:
+                last_reported_ts = None
+            else:
+                last_reported_ts = state.last_reported_timestamp
+        context = event.context
+        return States(
+            state=state_value,
+            entity_id=event.data["entity_id"],
+            attributes=None,
+            context_id=None,
+            context_id_bin=ulid_to_bytes_or_none(context.id),
+            context_user_id=None,
+            context_user_id_bin=uuid_hex_to_bytes_or_none(context.user_id),
+            context_parent_id=None,
+            context_parent_id_bin=ulid_to_bytes_or_none(context.parent_id),
+            origin_idx=event.origin.idx,
+            last_updated=None,
+            last_changed=None,
+            last_updated_ts=last_updated_ts,
+            last_changed_ts=last_changed_ts,
+            last_reported_ts=last_reported_ts,
+        )
 
     def to_native(self, validate_entity_id: bool = True) -> State | None:
         """Convert to an HA state object."""
diff --git a/homeassistant/core.py b/homeassistant/core.py
index c4392f62c52..cf5373ad8c2 100644
--- a/homeassistant/core.py
+++ b/homeassistant/core.py
@@ -1308,6 +1308,11 @@ class EventOrigin(enum.Enum):
         """Return the event."""
         return self.value
 
+    @cached_property
+    def idx(self) -> int:
+        """Return the index of the origin."""
+        return next((idx for idx, origin in enumerate(EventOrigin) if origin is self))
+
 
 class Event(Generic[_DataT]):
     """Representation of an event within the bus."""
diff --git a/tests/test_core.py b/tests/test_core.py
index 5e6b51cc39e..5f824f9e53a 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -920,6 +920,14 @@ def test_event_repr() -> None:
     )
 
 
+def test_event_origin_idx() -> None:
+    """Test the EventOrigin idx."""
+    assert ha.EventOrigin.remote is ha.EventOrigin.remote
+    assert ha.EventOrigin.local is ha.EventOrigin.local
+    assert ha.EventOrigin.local.idx == 0
+    assert ha.EventOrigin.remote.idx == 1
+
+
 def test_event_as_dict() -> None:
     """Test an Event as dictionary."""
     event_type = "some_type"