Use snapshot test helper in AO Smith (#115890)

tests/components/aosmith/test_sensor.py

@@ -1,50 +1,30 @@
 """Tests for the sensor platform of the A. O. Smith integration."""
 
+from collections.abc import AsyncGenerator
+from unittest.mock import patch
+
 import pytest
 from syrupy.assertion import SnapshotAssertion
 
+from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, snapshot_platform
 
 
-@pytest.mark.parametrize(
-    ("entity_id", "unique_id"),
-    [
-        (
-            "sensor.my_water_heater_hot_water_availability",
-            "hot_water_availability_junctionId",
-        ),
-        ("sensor.my_water_heater_energy_usage", "energy_usage_junctionId"),
-    ],
-)
-async def test_setup(
-    hass: HomeAssistant,
-    entity_registry: er.EntityRegistry,
-    init_integration: MockConfigEntry,
-    entity_id: str,
-    unique_id: str,
-) -> None:
-    """Test the setup of the sensor entities."""
-    entry = entity_registry.async_get(entity_id)
-    assert entry
-    assert entry.unique_id == unique_id
+@pytest.fixture(autouse=True)
+async def platforms() -> AsyncGenerator[list[str], None]:
+    """Return the platforms to be loaded for this test."""
+    with patch("homeassistant.components.aosmith.PLATFORMS", [Platform.SENSOR]):
+        yield
 
 
-@pytest.mark.parametrize(
-    ("entity_id"),
-    [
-        "sensor.my_water_heater_hot_water_availability",
-        "sensor.my_water_heater_energy_usage",
-    ],
-)
 async def test_state(
     hass: HomeAssistant,
     init_integration: MockConfigEntry,
     snapshot: SnapshotAssertion,
-    entity_id: str,
+    entity_registry: er.EntityRegistry,
 ) -> None:
     """Test the state of the sensor entities."""
-    state = hass.states.get(entity_id)
-    assert state == snapshot
+    await snapshot_platform(hass, entity_registry, snapshot, init_integration.entry_id)