Tweak usage prediction common control algorithm (home-assistant#152490)

Author: balloob

homeassistant/components/usage_prediction/common_control.py

@@ -18,6 +18,7 @@
 from homeassistant.components.recorder.util import session_scope
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import entity_registry as er
 from homeassistant.util import dt as dt_util
 from homeassistant.util.json import json_loads_object
 
@@ -41,8 +42,6 @@
     Platform.CAMERA,
     Platform.CLIMATE,
     Platform.COVER,
-    Platform.DATE,
-    Platform.DATETIME,
     Platform.FAN,
     Platform.HUMIDIFIER,
     Platform.IMAGE,
@@ -57,14 +56,9 @@
     Platform.SIREN,
     Platform.SWITCH,
     Platform.TEXT,
-    Platform.TIME,
-    Platform.TODO,
-    Platform.UPDATE,
     Platform.VACUUM,
     Platform.VALVE,
-    Platform.WAKE_WORD,
     Platform.WATER_HEATER,
-    Platform.WEATHER,
     # Helpers with own domain
     "counter",
     "group",
@@ -105,14 +99,17 @@ async def async_predict_common_control(
     """
     # Get the recorder instance to ensure it's ready
     recorder = get_instance(hass)
+    ent_reg = er.async_get(hass)
 
     # Execute the database operation in the recorder's executor
     return await recorder.async_add_executor_job(
-        _fetch_with_session, hass, _fetch_and_process_data, user_id
+        _fetch_with_session, hass, _fetch_and_process_data, ent_reg, user_id
     )
 
 
-def _fetch_and_process_data(session: Session, user_id: str) -> EntityUsagePredictions:
+def _fetch_and_process_data(
+    session: Session, ent_reg: er.EntityRegistry, user_id: str
+) -> EntityUsagePredictions:
     """Fetch and process service call events from the database."""
     # Prepare a dictionary to track results
     results: dict[str, Counter[str]] = {
@@ -198,6 +195,7 @@ def _fetch_and_process_data(session: Session, user_id: str) -> EntityUsagePredic
             entity_id
             for entity_id in entity_ids
             if entity_id.split(".")[0] in ALLOWED_DOMAINS
+            and ((entry := ent_reg.async_get(entity_id)) is None or not entry.hidden)
         ]
 
         if not entity_ids:

tests/components/usage_prediction/test_common_control.py

@@ -15,6 +15,7 @@
 from homeassistant.components.usage_prediction.models import EntityUsagePredictions
 from homeassistant.const import EVENT_CALL_SERVICE
 from homeassistant.core import Context, HomeAssistant
+from homeassistant.helpers import entity_registry as er
 
 from tests.components.recorder.common import async_wait_recording_done
 
@@ -135,6 +136,21 @@ async def test_multiple_entities_in_one_call(hass: HomeAssistant) -> None:
     """Test handling of service calls with multiple entity IDs."""
     user_id = str(uuid.uuid4())
 
+    ent_reg = er.async_get(hass)
+    ent_reg.async_get_or_create(
+        "light",
+        "test",
+        "living_room",
+        suggested_object_id="living_room",
+        hidden_by=er.RegistryEntryHider.USER,
+    )
+    ent_reg.async_get_or_create(
+        "light",
+        "test",
+        "kitchen",
+        suggested_object_id="kitchen",
+    )
+
     with freeze_time("2023-07-01 10:00:00+00:00"):  # Morning
         hass.bus.async_fire(
             EVENT_CALL_SERVICE,
@@ -159,8 +175,9 @@ async def test_multiple_entities_in_one_call(hass: HomeAssistant) -> None:
     with freeze_time("2023-07-02 10:00:00+00:00"):  # Next day, so events are recent
         results = await async_predict_common_control(hass, user_id)
 
-    # All three lights should be counted (10:00 UTC = 02:00 local = night)
-    assert results.night == ["light.living_room", "light.kitchen", "light.hallway"]
+    # Two lights should be counted (10:00 UTC = 02:00 local = night)
+    # Living room is hidden via entity registry
+    assert results.night == ["light.kitchen", "light.hallway"]
     assert results.morning == []
     assert results.afternoon == []
     assert results.evening == []