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

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

tests/components/labs/test_websocket_api.py

@@ -14,6 +14,8 @@ from homeassistant.components.labs import (
 )
 from homeassistant.core import HomeAssistant
 
+from . import assert_stored_labs_data
+
 from tests.common import MockUser
 from tests.typing import WebSocketGenerator
 
@@ -61,7 +63,9 @@ async def test_websocket_list_preview_features(
 
 
 async def test_websocket_update_preview_feature_enable(
-    hass: HomeAssistant, hass_ws_client: WebSocketGenerator
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    hass_storage: dict[str, Any],
 ) -> None:
     """Test enabling a preview feature via WebSocket."""
     # Load kitchen_sink integration
@@ -71,6 +75,8 @@ async def test_websocket_update_preview_feature_enable(
 
     client = await hass_ws_client(hass)
 
+    assert "core.labs" not in hass_storage
+
     # Track events
     events = []
 
@@ -103,6 +109,11 @@ async def test_websocket_update_preview_feature_enable(
     # Verify feature is now enabled
     assert async_is_preview_feature_enabled(hass, "kitchen_sink", "special_repair")
 
+    assert_stored_labs_data(
+        hass_storage,
+        [{"domain": "kitchen_sink", "preview_feature": "special_repair"}],
+    )
+
 
 async def test_websocket_update_preview_feature_disable(
     hass: HomeAssistant,
@@ -143,10 +154,16 @@ async def test_websocket_update_preview_feature_disable(
 
     # Verify feature is disabled
     assert not async_is_preview_feature_enabled(hass, "kitchen_sink", "special_repair")
+    assert_stored_labs_data(
+        hass_storage,
+        [],
+    )
 
 
 async def test_websocket_update_nonexistent_feature(
-    hass: HomeAssistant, hass_ws_client: WebSocketGenerator
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    hass_storage: dict[str, Any],
 ) -> None:
     """Test updating a preview feature that doesn't exist."""
     assert await async_setup(hass, {})
@@ -168,9 +185,13 @@ async def test_websocket_update_nonexistent_feature(
     assert msg["error"]["code"] == "not_found"
     assert "not found" in msg["error"]["message"].lower()
 
+    assert "core.labs" not in hass_storage
+
 
 async def test_websocket_update_unavailable_preview_feature(
-    hass: HomeAssistant, hass_ws_client: WebSocketGenerator
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    hass_storage: dict[str, Any],
 ) -> None:
     """Test updating a preview feature whose integration is not loaded still works."""
     # Don't load kitchen_sink integration
@@ -193,6 +214,11 @@ async def test_websocket_update_unavailable_preview_feature(
     assert msg["success"]
     assert msg["result"] is None
 
+    assert_stored_labs_data(
+        hass_storage,
+        [{"domain": "kitchen_sink", "preview_feature": "special_repair"}],
+    )
+
 
 @pytest.mark.parametrize(
     "command_type",
@@ -202,6 +228,7 @@ async def test_websocket_requires_admin(
     hass: HomeAssistant,
     hass_ws_client: WebSocketGenerator,
     hass_admin_user: MockUser,
+    hass_storage: dict[str, Any],
     command_type: str,
 ) -> None:
     """Test that websocket commands require admin privileges."""
@@ -230,6 +257,8 @@ async def test_websocket_requires_admin(
     assert not msg["success"]
     assert msg["error"]["code"] == "unauthorized"
 
+    assert "core.labs" not in hass_storage
+
 
 async def test_websocket_update_validates_enabled_parameter(
     hass: HomeAssistant, hass_ws_client: WebSocketGenerator
@@ -257,7 +286,9 @@ async def test_websocket_update_validates_enabled_parameter(
 
 
 async def test_storage_persists_preview_feature_across_calls(
-    hass: HomeAssistant, hass_ws_client: WebSocketGenerator
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    hass_storage: dict[str, Any],
 ) -> None:
     """Test that storage persists preview feature state across multiple calls."""
     hass.config.components.add("kitchen_sink")
@@ -266,6 +297,8 @@ async def test_storage_persists_preview_feature_across_calls(
 
     client = await hass_ws_client(hass)
 
+    assert "core.labs" not in hass_storage
+
     # Enable the preview feature
     await client.send_json_auto_id(
         {
@@ -278,6 +311,11 @@ async def test_storage_persists_preview_feature_across_calls(
     msg = await client.receive_json()
     assert msg["success"]
 
+    assert_stored_labs_data(
+        hass_storage,
+        [{"domain": "kitchen_sink", "preview_feature": "special_repair"}],
+    )
+
     # List preview features - should show enabled
     await client.send_json_auto_id({"type": "labs/list"})
     msg = await client.receive_json()
@@ -296,6 +334,11 @@ async def test_storage_persists_preview_feature_across_calls(
     msg = await client.receive_json()
     assert msg["success"]
 
+    assert_stored_labs_data(
+        hass_storage,
+        [],
+    )
+
     # List preview features - should show disabled
     await client.send_json_auto_id({"type": "labs/list"})
     msg = await client.receive_json()