From 35a1ecea272ae65b9e16f845f077389694d92806 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Mon, 3 Jun 2024 18:08:46 -0500
Subject: [PATCH] Speed up statistics_during_period websocket api (#118672)

---
 .../components/recorder/websocket_api.py      | 15 ++--
 .../components/recorder/test_websocket_api.py | 78 +++++++++++++++++++
 2 files changed, 85 insertions(+), 8 deletions(-)

diff --git a/homeassistant/components/recorder/websocket_api.py b/homeassistant/components/recorder/websocket_api.py
index 58c362df62e..b091343e5a4 100644
--- a/homeassistant/components/recorder/websocket_api.py
+++ b/homeassistant/components/recorder/websocket_api.py
@@ -160,14 +160,13 @@ def _ws_get_statistics_during_period(
         units,
         types,
     )
-    for statistic_id in result:
-        for item in result[statistic_id]:
-            if (start := item.get("start")) is not None:
-                item["start"] = int(start * 1000)
-            if (end := item.get("end")) is not None:
-                item["end"] = int(end * 1000)
-            if (last_reset := item.get("last_reset")) is not None:
-                item["last_reset"] = int(last_reset * 1000)
+    include_last_reset = "last_reset" in types
+    for statistic_rows in result.values():
+        for row in statistic_rows:
+            row["start"] = int(row["start"] * 1000)
+            row["end"] = int(row["end"] * 1000)
+            if include_last_reset and (last_reset := row["last_reset"]) is not None:
+                row["last_reset"] = int(last_reset * 1000)
     return json_bytes(messages.result_message(msg_id, result))
 
 
diff --git a/tests/components/recorder/test_websocket_api.py b/tests/components/recorder/test_websocket_api.py
index 9c8e0a9203a..3d35aafb2b3 100644
--- a/tests/components/recorder/test_websocket_api.py
+++ b/tests/components/recorder/test_websocket_api.py
@@ -3177,3 +3177,81 @@ async def test_adjust_sum_statistics_errors(
         stats = statistics_during_period(hass, zero, period="hour")
         assert stats != previous_stats
         previous_stats = stats
+
+
+async def test_import_statistics_with_last_reset(
+    recorder_mock: Recorder,
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test importing external statistics with last_reset can be fetched via websocket api."""
+    client = await hass_ws_client()
+
+    assert "Compiling statistics for" not in caplog.text
+    assert "Statistics already compiled" not in caplog.text
+
+    zero = dt_util.utcnow()
+    last_reset = dt_util.parse_datetime("2022-01-01T00:00:00+02:00")
+    period1 = zero.replace(minute=0, second=0, microsecond=0) + timedelta(hours=1)
+    period2 = zero.replace(minute=0, second=0, microsecond=0) + timedelta(hours=2)
+
+    external_statistics1 = {
+        "start": period1,
+        "last_reset": last_reset,
+        "state": 0,
+        "sum": 2,
+    }
+    external_statistics2 = {
+        "start": period2,
+        "last_reset": last_reset,
+        "state": 1,
+        "sum": 3,
+    }
+
+    external_metadata = {
+        "has_mean": False,
+        "has_sum": True,
+        "name": "Total imported energy",
+        "source": "test",
+        "statistic_id": "test:total_energy_import",
+        "unit_of_measurement": "kWh",
+    }
+
+    async_add_external_statistics(
+        hass, external_metadata, (external_statistics1, external_statistics2)
+    )
+    await async_wait_recording_done(hass)
+
+    client = await hass_ws_client()
+    await client.send_json_auto_id(
+        {
+            "type": "recorder/statistics_during_period",
+            "start_time": zero.isoformat(),
+            "end_time": (zero + timedelta(hours=48)).isoformat(),
+            "statistic_ids": ["test:total_energy_import"],
+            "period": "hour",
+            "types": ["change", "last_reset", "max", "mean", "min", "state", "sum"],
+        }
+    )
+    response = await client.receive_json()
+    assert response["result"] == {
+        "test:total_energy_import": [
+            {
+                "change": 2.0,
+                "end": (period1.timestamp() * 1000) + (3600 * 1000),
+                "last_reset": last_reset.timestamp() * 1000,
+                "start": period1.timestamp() * 1000,
+                "state": 0.0,
+                "sum": 2.0,
+            },
+            {
+                "change": 1.0,
+                "end": (period2.timestamp() * 1000 + (3600 * 1000)),
+                "last_reset": last_reset.timestamp() * 1000,
+                "start": period2.timestamp() * 1000,
+                "state": 1.0,
+                "sum": 3.0,
+            },
+        ]
+    }