From b780f10de88cecae2d61e2bf79e3cfe463b87953 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Mon, 10 Apr 2023 17:10:43 -1000
Subject: [PATCH] adjust

---
 .../components/recorder/history/modern.py     | 21 +++++++++++++++---
 .../components/recorder/models/state.py       | 22 +++++++++++--------
 tests/components/recorder/test_history.py     | 16 +++++++++-----
 3 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/homeassistant/components/recorder/history/modern.py b/homeassistant/components/recorder/history/modern.py
index 2fc7a816250..f03dd73a57d 100644
--- a/homeassistant/components/recorder/history/modern.py
+++ b/homeassistant/components/recorder/history/modern.py
@@ -651,7 +651,8 @@ def _sorted_states_to_dict(
     """
     field_map = _FIELD_MAP
     state_class: Callable[
-        [Row, dict[str, dict[str, Any]], float | None], State | dict[str, Any]
+        [Row, dict[str, dict[str, Any]], float | None, str, str, float | None],
+        State | dict[str, Any],
     ]
     if compressed_state_format:
         state_class = row_to_compressed_state
@@ -699,7 +700,14 @@ def _sorted_states_to_dict(
             or split_entity_id(entity_id)[0] in NEED_ATTRIBUTE_DOMAINS
         ):
             ent_results.extend(
-                state_class(db_state, attr_cache, start_time_ts, entity_id=entity_id)  # type: ignore[call-arg]
+                state_class(
+                    db_state,
+                    attr_cache,
+                    start_time_ts,
+                    entity_id,
+                    db_state[state_idx],
+                    db_state[last_updated_ts_idx],
+                )
                 for db_state in group
             )
             continue
@@ -713,7 +721,14 @@ def _sorted_states_to_dict(
                 continue
             prev_state = first_state[state_idx]
             ent_results.append(
-                state_class(first_state, attr_cache, start_time_ts, entity_id=entity_id)  # type: ignore[call-arg]
+                state_class(
+                    first_state,
+                    attr_cache,
+                    start_time_ts,
+                    entity_id,
+                    prev_state,  # type: ignore[arg-type]
+                    first_state[last_updated_ts_idx],
+                )
             )
 
         #
diff --git a/homeassistant/components/recorder/models/state.py b/homeassistant/components/recorder/models/state.py
index b6c6ba0897c..f31a5d95e43 100644
--- a/homeassistant/components/recorder/models/state.py
+++ b/homeassistant/components/recorder/models/state.py
@@ -52,16 +52,18 @@ class LazyState(State):
         row: Row,
         attr_cache: dict[str, dict[str, Any]],
         start_time_ts: float | None,
-        entity_id: str | None = None,
+        entity_id: str,
+        state: str,
+        last_updated_ts: float | None,
     ) -> None:
         """Init the lazy state."""
         self._row = row
-        self.entity_id = entity_id or self._row.entity_id
-        self.state = self._row.state or ""
+        self.entity_id = entity_id
+        self.state = state or ""
         self._attributes: dict[str, Any] | None = None
-        self._last_updated_ts: float | None = self._row.last_updated_ts or start_time_ts
+        self._last_updated_ts: float | None = last_updated_ts or start_time_ts
         self._last_changed_ts: float | None = (
-            getattr(self._row, "last_changed_ts", None) or self._last_updated_ts
+            getattr(self._row, "last_changed_ts", None) or last_updated_ts
         )
         self._context: Context | None = None
         self.attr_cache = attr_cache
@@ -137,14 +139,16 @@ def row_to_compressed_state(
     row: Row,
     attr_cache: dict[str, dict[str, Any]],
     start_time_ts: float | None,
-    entity_id: str | None = None,
+    entity_id: str,
+    state: str,
+    last_updated_ts: float | None,
 ) -> dict[str, Any]:
     """Convert a database row to a compressed state schema 31 and later."""
-    comp_state = {
-        COMPRESSED_STATE_STATE: row.state,
+    comp_state: dict[str, Any] = {
+        COMPRESSED_STATE_STATE: state,
         COMPRESSED_STATE_ATTRIBUTES: decode_attributes_from_row(row, attr_cache),
     }
-    row_last_updated_ts: float = row.last_updated_ts or start_time_ts  # type: ignore[assignment]
+    row_last_updated_ts: float = last_updated_ts or start_time_ts  # type: ignore[assignment]
     comp_state[COMPRESSED_STATE_LAST_UPDATED] = row_last_updated_ts
     if (
         (row_last_changed_ts := getattr(row, "last_changed_ts", None))
diff --git a/tests/components/recorder/test_history.py b/tests/components/recorder/test_history.py
index 5cd77447f1c..97498041725 100644
--- a/tests/components/recorder/test_history.py
+++ b/tests/components/recorder/test_history.py
@@ -55,14 +55,20 @@ async def _async_get_states(
     """Get states from the database."""
 
     def _get_states_with_session():
-        if get_instance(hass).schema_version < 31:
-            klass = LegacyLazyStatePreSchema31
-        else:
-            klass = LazyState
         with session_scope(hass=hass, read_only=True) as session:
             attr_cache = {}
+            pre_31_schema = get_instance(hass).schema_version < 31
             return [
-                klass(row, attr_cache, None)
+                LegacyLazyStatePreSchema31(row, attr_cache, None)
+                if pre_31_schema
+                else LazyState(
+                    row,
+                    attr_cache,
+                    None,
+                    row.entity_id,
+                    row.state,
+                    getattr(row, "last_updated_ts", None),
+                )
                 for row in legacy._get_rows_with_session(
                     hass,
                     session,