Fix use of storage helper in the labs integration (home-assistant#157249)

Co-authored-by: Copilot <[email protected]>

Author: emontnemery

homeassistant/components/labs/__init__.py

@@ -32,6 +32,7 @@
     LabPreviewFeature,
     LabsData,
     LabsStoreData,
+    NativeLabsStoreData,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -46,66 +47,12 @@
 ]
 
 
-class LabsStorage(Store[LabsStoreData]):
-    """Custom Store for Labs that converts between runtime and storage formats.
-
-    Runtime format: {"preview_feature_status": {(domain, preview_feature)}}
-    Storage format: {"preview_feature_status": [{"domain": str, "preview_feature": str}]}
-
-    Only enabled features are saved to storage - if stored, it's enabled.
-    """
-
-    async def _async_load_data(self) -> LabsStoreData | None:
-        """Load data and convert from storage format to runtime format."""
-        raw_data = await super()._async_load_data()
-        if raw_data is None:
-            return None
-
-        status_list = raw_data.get("preview_feature_status", [])
-
-        # Convert list of objects to runtime set - if stored, it's enabled
-        return {
-            "preview_feature_status": {
-                (item["domain"], item["preview_feature"]) for item in status_list
-            }
-        }
-
-    def _write_data(self, path: str, data: dict) -> None:
-        """Convert from runtime format to storage format and write.
-
-        Only saves enabled features - disabled is the default.
-        """
-        # Extract the actual data (has version/key wrapper)
-        actual_data = data.get("data", data)
-
-        # Check if this is Labs data (has preview_feature_status key)
-        if "preview_feature_status" not in actual_data:
-            # Not Labs data, write as-is
-            super()._write_data(path, data)
-            return
-
-        preview_status = actual_data["preview_feature_status"]
-
-        # Convert from runtime format (set of tuples) to storage format (list of dicts)
-        status_list = [
-            {"domain": domain, "preview_feature": preview_feature}
-            for domain, preview_feature in preview_status
-        ]
-
-        # Build the final data structure with converted format
-        data_copy = data.copy()
-        data_copy["data"] = {"preview_feature_status": status_list}
-
-        super()._write_data(path, data_copy)
-
-
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up the Labs component."""
-    store = LabsStorage(hass, STORAGE_VERSION, STORAGE_KEY, private=True)
-    data = await store.async_load()
-
-    if data is None:
-        data = {"preview_feature_status": set()}
+    store: Store[NativeLabsStoreData] = Store(
+        hass, STORAGE_VERSION, STORAGE_KEY, private=True
+    )
+    data = LabsStoreData.from_store_format(await store.async_load())
 
     # Scan ALL integrations for lab preview features (loaded or not)
     lab_preview_features = await _async_scan_all_preview_features(hass)
@@ -115,17 +62,17 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
         valid_keys = {
             (pf.domain, pf.preview_feature) for pf in lab_preview_features.values()
         }
-        stale_keys = data["preview_feature_status"] - valid_keys
+        stale_keys = data.preview_feature_status - valid_keys
 
         if stale_keys:
             _LOGGER.debug(
                 "Removing %d stale preview features: %s",
                 len(stale_keys),
                 stale_keys,
             )
-            data["preview_feature_status"] -= stale_keys
+            data.preview_feature_status -= stale_keys
 
-            await store.async_save(data)
+            await store.async_save(data.to_store_format())
 
     hass.data[LABS_DATA] = LabsData(
         store=store,
@@ -216,7 +163,7 @@ def async_is_preview_feature_enabled(
         return False
 
     labs_data = hass.data[LABS_DATA]
-    return (domain, preview_feature) in labs_data.data["preview_feature_status"]
+    return (domain, preview_feature) in labs_data.data.preview_feature_status
 
 
 @callback
@@ -265,7 +212,7 @@ def websocket_list_preview_features(
     preview_features: list[dict[str, Any]] = [
         preview_feature.to_dict(
             (preview_feature.domain, preview_feature.preview_feature)
-            in labs_data.data["preview_feature_status"]
+            in labs_data.data.preview_feature_status
         )
         for preview_feature in labs_data.preview_features.values()
         if preview_feature.domain in loaded_components
@@ -325,12 +272,12 @@ async def websocket_update_preview_feature(
 
     # Update storage (only store enabled features, remove if disabled)
     if enabled:
-        labs_data.data["preview_feature_status"].add((domain, preview_feature))
+        labs_data.data.preview_feature_status.add((domain, preview_feature))
     else:
-        labs_data.data["preview_feature_status"].discard((domain, preview_feature))
+        labs_data.data.preview_feature_status.discard((domain, preview_feature))
 
     # Save changes immediately
-    await labs_data.store.async_save(labs_data.data)
+    await labs_data.store.async_save(labs_data.data.to_store_format())
 
     # Fire event
     event_data: EventLabsUpdatedData = {

homeassistant/components/labs/const.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, TypedDict
+from typing import TYPE_CHECKING, Self, TypedDict
 
 from homeassistant.util.hass_dict import HassKey
 
@@ -62,14 +62,52 @@ def to_dict(self, enabled: bool) -> dict[str, str | bool | None]:
         }
 
 
-type LabsStoreData = dict[str, set[tuple[str, str]]]
+@dataclass(kw_only=True)
+class LabsStoreData:
+    """Storage data for Labs."""
+
+    preview_feature_status: set[tuple[str, str]]
+
+    @classmethod
+    def from_store_format(cls, data: NativeLabsStoreData | None) -> Self:
+        """Initialize from storage format."""
+        if data is None:
+            return cls(preview_feature_status=set())
+        return cls(
+            preview_feature_status={
+                (item["domain"], item["preview_feature"])
+                for item in data["preview_feature_status"]
+            }
+        )
+
+    def to_store_format(self) -> NativeLabsStoreData:
+        """Convert to storage format."""
+        return {
+            "preview_feature_status": [
+                {"domain": domain, "preview_feature": preview_feature}
+                for domain, preview_feature in self.preview_feature_status
+            ]
+        }
+
+
+class NativeLabsStoreData(TypedDict):
+    """Storage data for Labs."""
+
+    preview_feature_status: list[NativeLabsStoredFeature]
+
+
+class NativeLabsStoredFeature(TypedDict):
+    """A single preview feature entry in storage format."""
+
+    domain: str
+    preview_feature: str
 
 
 @dataclass
 class LabsData:
     """Storage class for Labs global data."""
 
-    store: Store[LabsStoreData]
+    store: Store[NativeLabsStoreData]
     data: LabsStoreData
     preview_features: dict[str, LabPreviewFeature] = field(default_factory=dict)
 

tests/components/labs/__init__.py

@@ -1 +1,18 @@
 """Tests for the Home Assistant Labs integration."""
+
+from typing import Any
+
+from pytest_unordered import unordered
+
+
+def assert_stored_labs_data(
+    hass_storage: dict[str, Any],
+    expected_data: list[dict[str, str]],
+) -> None:
+    """Assert that the storage has the expected enabled preview features."""
+    assert hass_storage["core.labs"] == {
+        "version": 1,
+        "minor_version": 1,
+        "key": "core.labs",
+        "data": {"preview_feature_status": unordered(expected_data)},
+    }

tests/components/labs/test_init.py

@@ -9,16 +9,16 @@
 
 from homeassistant.components.labs import (
     EVENT_LABS_UPDATED,
-    LabsStorage,
     async_is_preview_feature_enabled,
     async_listen,
 )
 from homeassistant.components.labs.const import DOMAIN, LABS_DATA, LabPreviewFeature
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.storage import Store
 from homeassistant.loader import Integration
 from homeassistant.setup import async_setup_component
 
+from . import assert_stored_labs_data
+
 
 async def test_async_setup(hass: HomeAssistant) -> None:
     """Test the Labs integration setup."""
@@ -91,7 +91,12 @@ async def test_async_is_preview_feature_enabled_when_disabled(
 
 
 @pytest.mark.parametrize(
-    ("features_to_store", "expected_enabled", "expected_cleaned"),
+    (
+        "features_to_store",
+        "expected_enabled",
+        "expected_cleaned",
+        "expected_cleaned_store",
+    ),
     [
         # Single stale feature cleanup
         (
@@ -101,6 +106,7 @@ async def test_async_is_preview_feature_enabled_when_disabled(
             ],
             [("kitchen_sink", "special_repair")],
             [("nonexistent_domain", "fake_feature")],
+            [{"domain": "kitchen_sink", "preview_feature": "special_repair"}],
         ),
         # Multiple stale features cleanup
         (
@@ -116,12 +122,14 @@ async def test_async_is_preview_feature_enabled_when_disabled(
                 ("stale_domain_2", "another_old"),
                 ("stale_domain_3", "yet_another"),
             ],
+            [{"domain": "kitchen_sink", "preview_feature": "special_repair"}],
         ),
         # All features cleaned (no integrations loaded)
         (
             [{"domain": "nonexistent", "preview_feature": "fake"}],
             [],
             [("nonexistent", "fake")],
+            [],
         ),
     ],
 )
@@ -131,6 +139,7 @@ async def test_storage_cleanup_stale_features(
     features_to_store: list[dict[str, str]],
     expected_enabled: list[tuple[str, str]],
     expected_cleaned: list[tuple[str, str]],
+    expected_cleaned_store: list[dict[str, str]],
 ) -> None:
     """Test that stale preview features are removed from storage on setup."""
     # Load kitchen_sink only if we expect any features to remain
@@ -155,6 +164,8 @@ async def test_storage_cleanup_stale_features(
     for domain, feature in expected_cleaned:
         assert not async_is_preview_feature_enabled(hass, domain, feature)
 
+    assert_stored_labs_data(hass_storage, expected_cleaned_store)
+
 
 @pytest.mark.parametrize(
     ("domain", "preview_feature", "expected"),
@@ -191,37 +202,6 @@ async def test_async_is_preview_feature_enabled(
     assert result is expected
 
 
-async def test_multiple_setups_idempotent(hass: HomeAssistant) -> None:
-    """Test that calling async_setup multiple times is safe."""
-    result1 = await async_setup_component(hass, DOMAIN, {})
-    assert result1 is True
-
-    result2 = await async_setup_component(hass, DOMAIN, {})
-    assert result2 is True
-
-    # Verify store is still accessible
-    assert LABS_DATA in hass.data
-
-
-async def test_storage_load_missing_preview_feature_status_key(
-    hass: HomeAssistant, hass_storage: dict[str, Any]
-) -> None:
-    """Test loading storage when preview_feature_status key is missing."""
-    # Storage data without preview_feature_status key
-    hass_storage["core.labs"] = {
-        "version": 1,
-        "minor_version": 1,
-        "key": "core.labs",
-        "data": {},  # Missing preview_feature_status
-    }
-
-    assert await async_setup_component(hass, DOMAIN, {})
-    await hass.async_block_till_done()
-
-    # Should initialize correctly - verify no feature is enabled
-    assert not async_is_preview_feature_enabled(hass, "kitchen_sink", "special_repair")
-
-
 async def test_preview_feature_full_key(hass: HomeAssistant) -> None:
     """Test that preview feature full_key property returns correct format."""
     feature = LabPreviewFeature(
@@ -276,24 +256,6 @@ async def test_preview_feature_to_dict_with_no_urls(hass: HomeAssistant) -> None
     }
 
 
-async def test_storage_load_returns_none_when_no_file(
-    hass: HomeAssistant,
-) -> None:
-    """Test storage load when no file exists (returns None)."""
-    # Create a storage instance but don't write any data
-    store = LabsStorage(hass, 1, "test_labs_none.json")
-
-    # Mock the parent Store's _async_load_data to return None
-    # This simulates the edge case where Store._async_load_data returns None
-    # This tests line 60: return None
-    async def mock_load_none():
-        return None
-
-    with patch.object(Store, "_async_load_data", new=mock_load_none):
-        result = await store.async_load()
-        assert result is None
-
-
 async def test_custom_integration_with_preview_features(
     hass: HomeAssistant,
 ) -> None: