Refactor tradfri tests (#110094)

* Refactor tradfri tests

* Refactor command store

* Fix fixture type annotations

* Fix test type errors

tests/components/tradfri/common.py

@@ -1,9 +1,12 @@
 """Common tools used for the Tradfri test suite."""
 from copy import deepcopy
+from dataclasses import dataclass
 from typing import Any
-from unittest.mock import Mock
 
+from pytradfri.command import Command
+from pytradfri.const import ATTR_ID
 from pytradfri.device import Device
+from pytradfri.gateway import Gateway
 
 from homeassistant.components import tradfri
 from homeassistant.core import HomeAssistant
@@ -13,7 +16,69 @@ from . import GATEWAY_ID
 from tests.common import MockConfigEntry
 
 
-async def setup_integration(hass):
+@dataclass
+class CommandStore:
+    """Store commands and command responses for the API."""
+
+    sent_commands: list[Command]
+    mock_responses: dict[str, Any]
+
+    def register_device(
+        self, gateway: Gateway, device_response: dict[str, Any]
+    ) -> None:
+        """Register device response."""
+        get_devices_command = gateway.get_devices()
+        self.register_response(get_devices_command, [device_response[ATTR_ID]])
+        get_device_command = gateway.get_device(device_response[ATTR_ID])
+        self.register_response(get_device_command, device_response)
+
+    def register_response(self, command: Command, response: Any) -> None:
+        """Register command response."""
+        self.mock_responses[command.path_str] = response
+
+    def process_command(self, command: Command) -> Any | None:
+        """Process command."""
+        response = self.mock_responses.get(command.path_str)
+        if response is None or command.process_result is None:
+            return None
+        return command.process_result(response)
+
+    async def trigger_observe_callback(
+        self,
+        hass: HomeAssistant,
+        device: Device,
+        new_device_state: dict[str, Any] | None = None,
+    ) -> None:
+        """Trigger the observe callback."""
+        observe_command = next(
+            (
+                command
+                for command in self.sent_commands
+                if command.path == device.path and command.observe
+            ),
+            None,
+        )
+        assert observe_command
+
+        device_path = "/".join(str(v) for v in device.path)
+        device_state = deepcopy(device.raw)
+
+        # Create a default observed state based on the sent commands.
+        for command in self.sent_commands:
+            if (data := command.data) is None or command.path_str != device_path:
+                continue
+            device_state = modify_state(device_state, data)
+
+        # Allow the test to override the default observed state.
+        if new_device_state is not None:
+            device_state = modify_state(device_state, new_device_state)
+
+        observe_command.process_result(device_state)
+
+        await hass.async_block_till_done()
+
+
+async def setup_integration(hass: HomeAssistant) -> MockConfigEntry:
     """Load the Tradfri integration with a mock gateway."""
     entry = MockConfigEntry(
         domain=tradfri.DOMAIN,
@@ -46,31 +111,3 @@ def modify_state(
             state[key] = value
 
     return state
-
-
-async def trigger_observe_callback(
-    hass: HomeAssistant,
-    mock_gateway: Mock,
-    device: Device,
-    new_device_state: dict[str, Any] | None = None,
-) -> None:
-    """Trigger the observe callback."""
-    observe_command = next(
-        (
-            command
-            for command in mock_gateway.mock_commands
-            if command.path == device.path and command.observe
-        ),
-        None,
-    )
-    assert observe_command
-
-    if new_device_state is not None:
-        mock_gateway.mock_responses.append(new_device_state)
-
-    device_state = deepcopy(device.raw)
-    new_state = mock_gateway.mock_responses[-1]
-    device_state = modify_state(device_state, new_state)
-    observe_command.process_result(device_state)
-
-    await hass.async_block_till_done()

tests/components/tradfri/conftest.py

@@ -1,120 +1,107 @@
 """Common tradfri test fixtures."""
 from __future__ import annotations
 
-from collections.abc import Generator
+from collections.abc import Callable, Generator
 import json
 from typing import Any
-from unittest.mock import MagicMock, Mock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from pytradfri.command import Command
+from pytradfri.const import ATTR_FIRMWARE_VERSION, ATTR_GATEWAY_ID
 from pytradfri.device import Device
-from pytradfri.device.air_purifier import AirPurifier
+from pytradfri.gateway import Gateway
-from pytradfri.device.blind import Blind
 
 from homeassistant.components.tradfri.const import DOMAIN
 
 from . import GATEWAY_ID, TRADFRI_PATH
+from .common import CommandStore
 
 from tests.common import load_fixture
 
 
 @pytest.fixture
-def mock_gateway_info():
+def mock_entry_setup() -> Generator[AsyncMock, None, None]:
-    """Mock get_gateway_info."""
-    with patch(f"{TRADFRI_PATH}.config_flow.get_gateway_info") as gateway_info:
-        yield gateway_info
-
-
-@pytest.fixture
-def mock_entry_setup():
     """Mock entry setup."""
     with patch(f"{TRADFRI_PATH}.async_setup_entry") as mock_setup:
         mock_setup.return_value = True
         yield mock_setup
 
 
-@pytest.fixture(name="mock_gateway")
+@pytest.fixture(name="mock_gateway", autouse=True)
-def mock_gateway_fixture():
+def mock_gateway_fixture(command_store: CommandStore) -> Gateway:
     """Mock a Tradfri gateway."""
-
+    gateway = Gateway()
-    def get_devices():
+    command_store.register_response(
-        """Return mock devices."""
+        gateway.get_gateway_info(),
-        return gateway.mock_devices
+        {ATTR_GATEWAY_ID: GATEWAY_ID, ATTR_FIRMWARE_VERSION: "1.2.1234"},
-
-    def get_groups():
-        """Return mock groups."""
-        return gateway.mock_groups
-
-    gateway_info = Mock(id=GATEWAY_ID, firmware_version="1.2.1234")
-
-    def get_gateway_info():
-        """Return mock gateway info."""
-        return gateway_info
-
-    gateway = Mock(
-        get_devices=get_devices,
-        get_groups=get_groups,
-        get_gateway_info=get_gateway_info,
-        mock_commands=[],
-        mock_devices=[],
-        mock_groups=[],
-        mock_responses=[],
     )
-    with patch(f"{TRADFRI_PATH}.Gateway", return_value=gateway), patch(
+    command_store.register_response(
-        f"{TRADFRI_PATH}.config_flow.Gateway", return_value=gateway
+        gateway.get_devices(),
-    ):
+        [],
-        yield gateway
+    )
+    return gateway
+
+
+@pytest.fixture(name="command_store", autouse=True)
+def command_store_fixture() -> CommandStore:
+    """Store commands and command responses for the API."""
+    return CommandStore([], {})
 
 
 @pytest.fixture(name="mock_api")
-def mock_api_fixture(mock_gateway):
+def mock_api_fixture(
+    command_store: CommandStore,
+) -> Callable[[Command | list[Command], float | None], Any | None]:
     """Mock api."""
 
-    async def api(command, timeout=None):
+    async def api(
+        command: Command | list[Command], timeout: float | None = None
+    ) -> Any | None:
         """Mock api function."""
-        # Store the data for "real" command objects.
+        if isinstance(command, list):
-        if hasattr(command, "_data") and not isinstance(command, Mock):
+            result = []
-            mock_gateway.mock_responses.append(command._data)
+            for cmd in command:
-            mock_gateway.mock_commands.append(command)
+                command_store.sent_commands.append(cmd)
-        return command
+                result.append(command_store.process_command(cmd))
+            return result
+
+        command_store.sent_commands.append(command)
+        return command_store.process_command(command)
 
     return api
 
 
-@pytest.fixture
+@pytest.fixture(autouse=True)
-def mock_api_factory(mock_api) -> Generator[MagicMock, None, None]:
+def mock_api_factory(
+    mock_api: Callable[[Command | list[Command], float | None], Any | None],
+) -> Generator[MagicMock, None, None]:
     """Mock pytradfri api factory."""
-    with patch(f"{TRADFRI_PATH}.APIFactory", autospec=True) as factory:
+    with patch(f"{TRADFRI_PATH}.APIFactory", autospec=True) as factory_class:
-        factory.init.return_value = factory.return_value
+        factory = factory_class.return_value
-        factory.return_value.request = mock_api
+        factory_class.init.return_value = factory
-        yield factory.return_value
+        factory.request = mock_api
+        yield factory
+
+
+@pytest.fixture
+def device(
+    command_store: CommandStore, mock_gateway: Gateway, request: pytest.FixtureRequest
+) -> Device:
+    """Return a device."""
+    device_response: dict[str, Any] = json.loads(request.getfixturevalue(request.param))
+    device = Device(device_response)
+    command_store.register_device(mock_gateway, device.raw)
+    return device
 
 
 @pytest.fixture(scope="session")
-def air_purifier_response() -> dict[str, Any]:
+def air_purifier() -> str:
     """Return an air purifier response."""
-    return json.loads(load_fixture("air_purifier.json", DOMAIN))
+    return load_fixture("air_purifier.json", DOMAIN)
-
-
-@pytest.fixture
-def air_purifier(air_purifier_response: dict[str, Any]) -> AirPurifier:
-    """Return air purifier."""
-    device = Device(air_purifier_response)
-    air_purifier_control = device.air_purifier_control
-    assert air_purifier_control
-    return air_purifier_control.air_purifiers[0]
 
 
 @pytest.fixture(scope="session")
-def blind_response() -> dict[str, Any]:
+def blind() -> str:
     """Return a blind response."""
-    return json.loads(load_fixture("blind.json", DOMAIN))
+    return load_fixture("blind.json", DOMAIN)
-
-
-@pytest.fixture
-def blind(blind_response: dict[str, Any]) -> Blind:
-    """Return blind."""
-    device = Device(blind_response)
-    blind_control = device.blind_control
-    assert blind_control
-    return blind_control.blinds[0]

tests/components/tradfri/fixtures/bulb_cws.json

@@ -0,0 +1,28 @@
+{
+  "3": {
+    "0": "IKEA of Sweden",
+    "1": "TRADFRI bulb E27 CWS opal 600lm",
+    "2": "",
+    "3": "1.3.002",
+    "6": 1
+  },
+  "3311": [
+    {
+      "5706": "f1e0b5",
+      "5707": 5427,
+      "5708": 42596,
+      "5709": 30015,
+      "5710": 26870,
+      "5850": 1,
+      "5851": 250,
+      "9003": 0
+    }
+  ],
+  "5750": 2,
+  "9001": "Test CWS",
+  "9002": 1509924799,
+  "9003": 65541,
+  "9019": 1,
+  "9020": 1510011206,
+  "9054": 0
+}

tests/components/tradfri/fixtures/bulb_w.json

@@ -0,0 +1,17 @@
+{
+  "3": {
+    "0": "IKEA of Sweden",
+    "1": "TRADFRI bulb E27 W opal 1000lm",
+    "2": "",
+    "3": "1.2.214",
+    "6": 1
+  },
+  "3311": [{ "5850": 1, "5851": 250, "9003": 0 }],
+  "5750": 2,
+  "9001": "Test W",
+  "9002": 1509923551,
+  "9003": 65537,
+  "9019": 1,
+  "9020": 1510009959,
+  "9054": 0
+}

tests/components/tradfri/fixtures/bulb_ws.json

@@ -0,0 +1,27 @@
+{
+  "3": {
+    "0": "IKEA of Sweden",
+    "1": "TRADFRI bulb E27 WS opal 980lm",
+    "2": "",
+    "3": "1.2.217",
+    "6": 1
+  },
+  "3311": [
+    {
+      "5706": "0",
+      "5709": 31103,
+      "5710": 27007,
+      "5711": 400,
+      "5850": 1,
+      "5851": 250,
+      "9003": 0
+    }
+  ],
+  "5750": 2,
+  "9001": "Test WS",
+  "9002": 1509923713,
+  "9003": 65539,
+  "9019": 1,
+  "9020": 1510010121,
+  "9054": 0
+}

tests/components/tradfri/test_cover.py

@@ -2,29 +2,26 @@
 from __future__ import annotations
 
 from typing import Any
-from unittest.mock import MagicMock, Mock
 
 import pytest
 from pytradfri.const import ATTR_REACHABLE_STATE
-from pytradfri.device.blind import Blind
+from pytradfri.device import Device
 
 from homeassistant.components.cover import ATTR_CURRENT_POSITION, DOMAIN as COVER_DOMAIN
 from homeassistant.const import STATE_CLOSED, STATE_OPEN, STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
 
-from .common import setup_integration, trigger_observe_callback
+from .common import CommandStore, setup_integration
 
 
+@pytest.mark.parametrize("device", ["blind"], indirect=True)
 async def test_cover_available(
     hass: HomeAssistant,
-    mock_gateway: Mock,
+    command_store: CommandStore,
-    mock_api_factory: MagicMock,
+    device: Device,
-    blind: Blind,
 ) -> None:
     """Test cover available property."""
     entity_id = "cover.test"
-    device = blind.device
-    mock_gateway.mock_devices.append(device)
     await setup_integration(hass)
 
     state = hass.states.get(entity_id)
@@ -33,8 +30,8 @@ async def test_cover_available(
     assert state.attributes[ATTR_CURRENT_POSITION] == 60
     assert state.attributes["model"] == "FYRTUR block-out roller blind"
 
-    await trigger_observe_callback(
+    await command_store.trigger_observe_callback(
-        hass, mock_gateway, device, {ATTR_REACHABLE_STATE: 0}
+        hass, device, {ATTR_REACHABLE_STATE: 0}
     )
 
     state = hass.states.get(entity_id)
@@ -42,6 +39,7 @@ async def test_cover_available(
     assert state.state == STATE_UNAVAILABLE
 
 
+@pytest.mark.parametrize("device", ["blind"], indirect=True)
 @pytest.mark.parametrize(
     ("service", "service_data", "expected_state", "expected_position"),
     [
@@ -54,9 +52,8 @@ async def test_cover_available(
 )
 async def test_cover_services(
     hass: HomeAssistant,
-    mock_gateway: Mock,
+    command_store: CommandStore,
-    mock_api_factory: MagicMock,
+    device: Device,
-    blind: Blind,
     service: str,
     service_data: dict[str, Any],
     expected_state: str,
@@ -64,8 +61,6 @@ async def test_cover_services(
 ) -> None:
     """Test cover services."""
     entity_id = "cover.test"
-    device = blind.device
-    mock_gateway.mock_devices.append(device)
     await setup_integration(hass)
 
     state = hass.states.get(entity_id)
@@ -81,7 +76,7 @@ async def test_cover_services(
     )
     await hass.async_block_till_done()
 
-    await trigger_observe_callback(hass, mock_gateway, device)
+    await command_store.trigger_observe_callback(hass, device)
 
     state = hass.states.get(entity_id)
     assert state

tests/components/tradfri/test_diagnostics.py

@@ -1,9 +1,8 @@
 """Tests for Tradfri diagnostics."""
 from __future__ import annotations
 
-from unittest.mock import MagicMock, Mock
+import pytest
-
+from pytradfri.device import Device
-from pytradfri.device.air_purifier import AirPurifier
 
 from homeassistant.core import HomeAssistant
 
@@ -13,16 +12,13 @@ from tests.components.diagnostics import get_diagnostics_for_config_entry
 from tests.typing import ClientSessionGenerator
 
 
+@pytest.mark.parametrize("device", ["air_purifier"], indirect=True)
 async def test_diagnostics(
     hass: HomeAssistant,
     hass_client: ClientSessionGenerator,
-    mock_gateway: Mock,
+    device: Device,
-    mock_api_factory: MagicMock,
-    air_purifier: AirPurifier,
 ) -> None:
     """Test diagnostics for config entry."""
-    device = air_purifier.device
-    mock_gateway.mock_devices.append(device)
     config_entry = await setup_integration(hass)
 
     result = await get_diagnostics_for_config_entry(hass, hass_client, config_entry)

tests/components/tradfri/test_fan.py

@@ -2,7 +2,6 @@
 from __future__ import annotations
 
 from typing import Any
-from unittest.mock import MagicMock, Mock
 
 import pytest
 from pytradfri.const import (
@@ -11,7 +10,7 @@ from pytradfri.const import (
     ATTR_REACHABLE_STATE,
     ROOT_AIR_PURIFIER,
 )
-from pytradfri.device.air_purifier import AirPurifier
+from pytradfri.device import Device
 
 from homeassistant.components.fan import (
     ATTR_PERCENTAGE,
@@ -32,19 +31,17 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant
 
-from .common import setup_integration, trigger_observe_callback
+from .common import CommandStore, setup_integration
 
 
+@pytest.mark.parametrize("device", ["air_purifier"], indirect=True)
 async def test_fan_available(
     hass: HomeAssistant,
-    mock_gateway: Mock,
+    command_store: CommandStore,
-    mock_api_factory: MagicMock,
+    device: Device,
-    air_purifier: AirPurifier,
 ) -> None:
     """Test fan available property."""
     entity_id = "fan.test"
-    device = air_purifier.device
-    mock_gateway.mock_devices.append(device)
     await setup_integration(hass)
 
     state = hass.states.get(entity_id)
@@ -56,8 +53,8 @@ async def test_fan_available(
     assert state.attributes[ATTR_PRESET_MODE] is None
     assert state.attributes[ATTR_SUPPORTED_FEATURES] == 9
 
-    await trigger_observe_callback(
+    await command_store.trigger_observe_callback(
-        hass, mock_gateway, device, {ATTR_REACHABLE_STATE: 0}
+        hass, device, {ATTR_REACHABLE_STATE: 0}
     )
 
     state = hass.states.get(entity_id)
@@ -65,6 +62,7 @@ async def test_fan_available(
     assert state.state == STATE_UNAVAILABLE
 
 
+@pytest.mark.parametrize("device", ["air_purifier"], indirect=True)
 @pytest.mark.parametrize(
     (
         "service",
@@ -153,9 +151,8 @@ async def test_fan_available(
 )
 async def test_services(
     hass: HomeAssistant,
-    mock_gateway: Mock,
+    command_store: CommandStore,
-    mock_api_factory: MagicMock,
+    device: Device,
-    air_purifier: AirPurifier,
     service: str,
     service_data: dict[str, Any],
     device_state: dict[str, Any],
@@ -165,8 +162,6 @@ async def test_services(
 ) -> None:
     """Test fan services."""
     entity_id = "fan.test"
-    device = air_purifier.device
-    mock_gateway.mock_devices.append(device)
     await setup_integration(hass)
 
     state = hass.states.get(entity_id)
@@ -186,9 +181,8 @@ async def test_services(
     )
     await hass.async_block_till_done()
 
-    await trigger_observe_callback(
+    await command_store.trigger_observe_callback(
         hass,
-        mock_gateway,
         device,
         {ROOT_AIR_PURIFIER: [device_state]},
     )

tests/components/tradfri/test_init.py

@@ -1,5 +1,5 @@
 """Tests for Tradfri setup."""
-from unittest.mock import patch
+from unittest.mock import MagicMock
 
 from homeassistant.components import tradfri
 from homeassistant.core import HomeAssistant
@@ -10,9 +10,11 @@ from . import GATEWAY_ID
 from tests.common import MockConfigEntry
 
 
-async def test_entry_setup_unload(hass: HomeAssistant, mock_api_factory) -> None:
+async def test_entry_setup_unload(
+    hass: HomeAssistant, device_registry: dr.DeviceRegistry, mock_api_factory: MagicMock
+) -> None:
     """Test config entry setup and unload."""
-    entry = MockConfigEntry(
+    config_entry = MockConfigEntry(
         domain=tradfri.DOMAIN,
         data={
             tradfri.CONF_HOST: "mock-host",
@@ -22,70 +24,69 @@ async def test_entry_setup_unload(hass: HomeAssistant, mock_api_factory) -> None
         },
     )
 
-    entry.add_to_hass(hass)
+    config_entry.add_to_hass(hass)
-    with patch.object(
+    await hass.config_entries.async_setup(config_entry.entry_id)
-        hass.config_entries, "async_forward_entry_setup", return_value=True
-    ) as setup:
-        await hass.config_entries.async_setup(entry.entry_id)
-        await hass.async_block_till_done()
-        assert setup.call_count == len(tradfri.PLATFORMS)
-
-    dev_reg = dr.async_get(hass)
-    dev_entries = dr.async_entries_for_config_entry(dev_reg, entry.entry_id)
-
-    assert dev_entries
-    dev_entry = dev_entries[0]
-    assert dev_entry.identifiers == {
-        (tradfri.DOMAIN, entry.data[tradfri.CONF_GATEWAY_ID])
-    }
-    assert dev_entry.manufacturer == "IKEA of Sweden"
-    assert dev_entry.name == "Gateway"
-    assert dev_entry.model == "E1526"
-
-    with patch.object(
-        hass.config_entries, "async_forward_entry_unload", return_value=True
-    ) as unload:
-        assert await hass.config_entries.async_unload(entry.entry_id)
-        await hass.async_block_till_done()
-        assert unload.call_count == len(tradfri.PLATFORMS)
-        assert mock_api_factory.shutdown.call_count == 1
-
-
-async def test_remove_stale_devices(hass: HomeAssistant, mock_api_factory) -> None:
-    """Test remove stale device registry entries."""
-    entry = MockConfigEntry(
-        domain=tradfri.DOMAIN,
-        data={
-            tradfri.CONF_HOST: "mock-host",
-            tradfri.CONF_IDENTITY: "mock-identity",
-            tradfri.CONF_KEY: "mock-key",
-            tradfri.CONF_GATEWAY_ID: GATEWAY_ID,
-        },
-    )
-
-    entry.add_to_hass(hass)
-    dev_reg = dr.async_get(hass)
-    dev_reg.async_get_or_create(
-        config_entry_id=entry.entry_id,
-        identifiers={(tradfri.DOMAIN, "stale_device_id")},
-    )
-    dev_entries = dr.async_entries_for_config_entry(dev_reg, entry.entry_id)
-
-    assert len(dev_entries) == 1
-    dev_entry = dev_entries[0]
-    assert dev_entry.identifiers == {(tradfri.DOMAIN, "stale_device_id")}
-
-    await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
-    dev_entries = dr.async_entries_for_config_entry(dev_reg, entry.entry_id)
+    device_entries = dr.async_entries_for_config_entry(
+        device_registry, config_entry.entry_id
+    )
+
+    assert device_entries
+    device_entry = device_entries[0]
+    assert device_entry.identifiers == {
+        (tradfri.DOMAIN, config_entry.data[tradfri.CONF_GATEWAY_ID])
+    }
+    assert device_entry.manufacturer == "IKEA of Sweden"
+    assert device_entry.name == "Gateway"
+    assert device_entry.model == "E1526"
+
+    assert await hass.config_entries.async_unload(config_entry.entry_id)
+    await hass.async_block_till_done()
+    assert mock_api_factory.shutdown.call_count == 1
+
+
+async def test_remove_stale_devices(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+) -> None:
+    """Test remove stale device registry entries."""
+    config_entry = MockConfigEntry(
+        domain=tradfri.DOMAIN,
+        data={
+            tradfri.CONF_HOST: "mock-host",
+            tradfri.CONF_IDENTITY: "mock-identity",
+            tradfri.CONF_KEY: "mock-key",
+            tradfri.CONF_GATEWAY_ID: GATEWAY_ID,
+        },
+    )
+
+    config_entry.add_to_hass(hass)
+    device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        identifiers={(tradfri.DOMAIN, "stale_device_id")},
+    )
+    device_entries = dr.async_entries_for_config_entry(
+        device_registry, config_entry.entry_id
+    )
+
+    assert len(device_entries) == 1
+    device_entry = device_entries[0]
+    assert device_entry.identifiers == {(tradfri.DOMAIN, "stale_device_id")}
+
+    await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    device_entries = dr.async_entries_for_config_entry(
+        device_registry, config_entry.entry_id
+    )
 
     # Check that only the gateway device entry remains.
-    assert len(dev_entries) == 1
+    assert len(device_entries) == 1
-    dev_entry = dev_entries[0]
+    device_entry = device_entries[0]
-    assert dev_entry.identifiers == {
+    assert device_entry.identifiers == {
-        (tradfri.DOMAIN, entry.data[tradfri.CONF_GATEWAY_ID])
+        (tradfri.DOMAIN, config_entry.data[tradfri.CONF_GATEWAY_ID])
     }
-    assert dev_entry.manufacturer == "IKEA of Sweden"
+    assert device_entry.manufacturer == "IKEA of Sweden"
-    assert dev_entry.name == "Gateway"
+    assert device_entry.name == "Gateway"
-    assert dev_entry.model == "E1526"
+    assert device_entry.model == "E1526"

tests/components/tradfri/test_light.py

@@ -1,310 +1,309 @@
 """Tradfri lights platform tests."""
-from copy import deepcopy
+from typing import Any
-from unittest.mock import MagicMock, Mock, PropertyMock, patch
 
 import pytest
+from pytradfri.const import ATTR_DEVICE_STATE, ATTR_LIGHT_CONTROL, ATTR_REACHABLE_STATE
 from pytradfri.device import Device
-from pytradfri.device.light import Light
-from pytradfri.device.light_control import LightControl
 
+from homeassistant.components.light import (
+    ATTR_BRIGHTNESS,
+    ATTR_COLOR_MODE,
+    ATTR_COLOR_TEMP,
+    ATTR_HS_COLOR,
+    ATTR_MAX_MIREDS,
+    ATTR_MIN_MIREDS,
+    ATTR_SUPPORTED_COLOR_MODES,
+    DOMAIN as LIGHT_DOMAIN,
+    ColorMode,
+)
+from homeassistant.components.tradfri.const import DOMAIN
+from homeassistant.const import (
+    SERVICE_TURN_OFF,
+    SERVICE_TURN_ON,
+    STATE_OFF,
+    STATE_ON,
+    STATE_UNAVAILABLE,
+)
 from homeassistant.core import HomeAssistant
 
-from .common import setup_integration
+from .common import CommandStore, setup_integration
 
-DEFAULT_TEST_FEATURES = {
+from tests.common import load_fixture
-    "can_set_dimmer": False,
+
-    "can_set_color": False,
+
-    "can_set_temp": False,
+@pytest.fixture(scope="module")
-}
+def bulb_w() -> str:
-# [
+    """Return a bulb W response."""
-#     {bulb features},
+    return load_fixture("bulb_w.json", DOMAIN)
-#     {turn_on arguments},
+
-#     {expected result}
+
-# ]
+@pytest.fixture(scope="module")
-TURN_ON_TEST_CASES = [
+def bulb_ws() -> str:
-    # Turn On
+    """Return a bulb WS response."""
-    [{}, {}, {"state": "on"}],
+    return load_fixture("bulb_ws.json", DOMAIN)
-    # Brightness > 0
+
-    [{"can_set_dimmer": True}, {"brightness": 100}, {"state": "on", "brightness": 100}],
+
-    # Brightness == 1
+@pytest.fixture(scope="module")
-    [{"can_set_dimmer": True}, {"brightness": 1}, {"brightness": 1}],
+def bulb_cws() -> str:
-    # Brightness > 254
+    """Return a bulb CWS response."""
-    [{"can_set_dimmer": True}, {"brightness": 1000}, {"brightness": 254}],
+    return load_fixture("bulb_cws.json", DOMAIN)
-    # color_temp
+
-    [{"can_set_temp": True}, {"color_temp": 250}, {"color_temp": 250}],
+
-    # color_temp < 250
+@pytest.mark.parametrize(
-    [{"can_set_temp": True}, {"color_temp": 1}, {"color_temp": 250}],
+    ("device", "entity_id", "state_attributes"),
-    # color_temp > 454
-    [{"can_set_temp": True}, {"color_temp": 1000}, {"color_temp": 454}],
-    # hs color
     [
-        {"can_set_color": True},
+        (
-        {"hs_color": [300, 100]},
+            "bulb_w",
-        {"state": "on", "hs_color": [300, 100]},
+            "light.test_w",
+            {
+                ATTR_BRIGHTNESS: 250,
+                ATTR_SUPPORTED_COLOR_MODES: [ColorMode.BRIGHTNESS],
+                ATTR_COLOR_MODE: ColorMode.BRIGHTNESS,
+            },
+        ),
+        (
+            "bulb_ws",
+            "light.test_ws",
+            {
+                ATTR_BRIGHTNESS: 250,
+                ATTR_COLOR_TEMP: 400,
+                ATTR_MIN_MIREDS: 250,
+                ATTR_MAX_MIREDS: 454,
+                ATTR_SUPPORTED_COLOR_MODES: [ColorMode.COLOR_TEMP],
+                ATTR_COLOR_MODE: ColorMode.COLOR_TEMP,
+            },
+        ),
+        (
+            "bulb_cws",
+            "light.test_cws",
+            {
+                ATTR_BRIGHTNESS: 250,
+                ATTR_HS_COLOR: (29.812, 65.252),
+                ATTR_SUPPORTED_COLOR_MODES: [ColorMode.HS],
+                ATTR_COLOR_MODE: ColorMode.HS,
+            },
+        ),
     ],
-    # ct + brightness
+    indirect=["device"],
-    [
+)
-        {"can_set_dimmer": True, "can_set_temp": True},
+async def test_light_state(
-        {"color_temp": 250, "brightness": 200},
+    hass: HomeAssistant,
-        {"state": "on", "color_temp": 250, "brightness": 200},
+    device: Device,
-    ],
+    entity_id: str,
-    # ct + brightness (no temp support)
+    state_attributes: dict[str, Any],
-    [
-        {"can_set_dimmer": True, "can_set_temp": False, "can_set_color": True},
-        {"color_temp": 250, "brightness": 200},
-        {"state": "on", "hs_color": [26.807, 34.869], "brightness": 200},
-    ],
-    # ct + brightness (no temp or color support)
-    [
-        {"can_set_dimmer": True, "can_set_temp": False, "can_set_color": False},
-        {"color_temp": 250, "brightness": 200},
-        {"state": "on", "brightness": 200},
-    ],
-    # hs + brightness
-    [
-        {"can_set_dimmer": True, "can_set_color": True},
-        {"hs_color": [300, 100], "brightness": 200},
-        {"state": "on", "hs_color": [300, 100], "brightness": 200},
-    ],
-]
-
-# Result of transition is not tested, but data is passed to turn on service.
-TRANSITION_CASES_FOR_TESTS = [None, 0, 1]
-
-
-@pytest.fixture(autouse=True, scope="module")
-def setup():
-    """Set up patches for pytradfri methods."""
-    p_1 = patch(
-        "pytradfri.device.LightControl.raw",
-        new_callable=PropertyMock,
-        return_value=[{"mock": "mock"}],
-    )
-    p_2 = patch("pytradfri.device.LightControl.lights")
-    p_1.start()
-    p_2.start()
-
-    yield
-
-    p_1.stop()
-    p_2.stop()
-
-
-async def generate_psk(self, code):
-    """Mock psk."""
-    return "mock"
-
-
-def mock_light(test_features=None, test_state=None, light_number=0):
-    """Mock a tradfri light."""
-    if test_features is None:
-        test_features = {}
-    if test_state is None:
-        test_state = {}
-    mock_light_data = Mock(**test_state)
-
-    dev_info_mock = MagicMock()
-    dev_info_mock.manufacturer = "manufacturer"
-    dev_info_mock.model_number = "model"
-    dev_info_mock.firmware_version = "1.2.3"
-    _mock_light = Mock(
-        id=f"mock-light-id-{light_number}",
-        reachable=True,
-        observe=Mock(),
-        device_info=dev_info_mock,
-        has_light_control=True,
-        has_socket_control=False,
-        has_blind_control=False,
-        has_signal_repeater_control=False,
-        has_air_purifier_control=False,
-    )
-    _mock_light.name = f"tradfri_light_{light_number}"
-
-    # Set supported features for the light.
-    features = {**DEFAULT_TEST_FEATURES, **test_features}
-    light_control = LightControl(_mock_light)
-    for attr, value in features.items():
-        setattr(light_control, attr, value)
-    # Store the initial state.
-    setattr(light_control, "lights", [mock_light_data])
-    _mock_light.light_control = light_control
-    return _mock_light
-
-
-async def test_light(hass: HomeAssistant, mock_gateway, mock_api_factory) -> None:
-    """Test that lights are correctly added."""
-    features = {"can_set_dimmer": True, "can_set_color": True, "can_set_temp": True}
-
-    state = {
-        "state": True,
-        "dimmer": 100,
-        "color_temp": 250,
-        "hsb_xy_color": (100, 100, 100, 100, 100),
-    }
-
-    mock_gateway.mock_devices.append(
-        mock_light(test_features=features, test_state=state)
-    )
-    await setup_integration(hass)
-
-    lamp_1 = hass.states.get("light.tradfri_light_0")
-    assert lamp_1 is not None
-    assert lamp_1.state == "on"
-    assert lamp_1.attributes["brightness"] == 100
-    assert lamp_1.attributes["hs_color"] == (0.549, 0.153)
-
-
-async def test_light_observed(
-    hass: HomeAssistant, mock_gateway, mock_api_factory
 ) -> None:
-    """Test that lights are correctly observed."""
+    """Test light state."""
-    light = mock_light()
-    mock_gateway.mock_devices.append(light)
     await setup_integration(hass)
-    assert len(light.observe.mock_calls) > 0
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == STATE_ON
+    for key, value in state_attributes.items():
+        assert state.attributes[key] == value
 
 
+@pytest.mark.parametrize("device", ["bulb_w"], indirect=True)
 async def test_light_available(
-    hass: HomeAssistant, mock_gateway, mock_api_factory
+    hass: HomeAssistant,
+    command_store: CommandStore,
+    device: Device,
 ) -> None:
     """Test light available property."""
-    light = mock_light({"state": True}, light_number=1)
+    entity_id = "light.test_w"
-    light.reachable = True
-
-    light2 = mock_light({"state": True}, light_number=2)
-    light2.reachable = False
-
-    mock_gateway.mock_devices.append(light)
-    mock_gateway.mock_devices.append(light2)
     await setup_integration(hass)
 
-    assert hass.states.get("light.tradfri_light_1").state == "on"
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == STATE_ON
 
-    assert hass.states.get("light.tradfri_light_2").state == "unavailable"
+    await command_store.trigger_observe_callback(
+        hass, device, {ATTR_REACHABLE_STATE: 0}
+    )
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == STATE_UNAVAILABLE
 
 
-def create_all_turn_on_cases():
+@pytest.mark.parametrize(
-    """Create all turn on test cases."""
+    "transition",
-    # Combine TURN_ON_TEST_CASES and TRANSITION_CASES_FOR_TESTS
+    [{}, {"transition": 0}, {"transition": 1}],
-    all_turn_on_test_cases = [
+    ids=["transition_none", "transition_0", "transition_1"],
-        ["test_features", "test_data", "expected_result", "device_id"],
+)
-        [],
+@pytest.mark.parametrize(
-    ]
+    ("device", "entity_id", "service_data", "state_attributes"),
-    index = 1
+    [
-    for test_case in TURN_ON_TEST_CASES:
+        # turn_on
-        for trans in TRANSITION_CASES_FOR_TESTS:
+        (
-            case = deepcopy(test_case)
+            "bulb_w",
-            if trans is not None:
+            "light.test_w",
-                case[1]["transition"] = trans
+            {},
-            case.append(index)
+            {},
-            index += 1
+        ),
-            all_turn_on_test_cases[1].append(case)
+        # brightness > 0
-
+        (
-    return all_turn_on_test_cases
+            "bulb_w",
-
+            "light.test_w",
-
+            {"brightness": 100},
-@pytest.mark.parametrize(*create_all_turn_on_cases())
+            {"brightness": 100},
+        ),
+        # brightness == 1
+        (
+            "bulb_w",
+            "light.test_w",
+            {"brightness": 1},
+            {"brightness": 1},
+        ),
+        # brightness > 254
+        (
+            "bulb_w",
+            "light.test_w",
+            {"brightness": 1000},
+            {"brightness": 254},
+        ),
+        # color_temp
+        (
+            "bulb_ws",
+            "light.test_ws",
+            {"color_temp": 250},
+            {"color_temp": 250},
+        ),
+        # color_temp < 250
+        (
+            "bulb_ws",
+            "light.test_ws",
+            {"color_temp": 1},
+            {"color_temp": 250},
+        ),
+        # color_temp > 454
+        (
+            "bulb_ws",
+            "light.test_ws",
+            {"color_temp": 1000},
+            {"color_temp": 454},
+        ),
+        # hs_color
+        (
+            "bulb_cws",
+            "light.test_cws",
+            {"hs_color": [300, 100]},
+            {"hs_color": [300, 100]},
+        ),
+        # ct + brightness
+        (
+            "bulb_ws",
+            "light.test_ws",
+            {"color_temp": 250, "brightness": 200},
+            {"color_temp": 250, "brightness": 200},
+        ),
+        # ct + brightness (no temp support)
+        (
+            "bulb_cws",
+            "light.test_cws",
+            {"color_temp": 250, "brightness": 200},
+            {"hs_color": [26.807, 34.869], "brightness": 200},
+        ),
+        # ct + brightness (no temp or color support)
+        (
+            "bulb_w",
+            "light.test_w",
+            {"color_temp": 250, "brightness": 200},
+            {"brightness": 200},
+        ),
+        # hs + brightness
+        (
+            "bulb_cws",
+            "light.test_cws",
+            {"hs_color": [300, 100], "brightness": 200},
+            {"hs_color": [300, 100], "brightness": 200},
+        ),
+    ],
+    indirect=["device"],
+    ids=[
+        "turn_on",
+        "brightness > 0",
+        "brightness == 1",
+        "brightness > 254",
+        "color_temp",
+        "color_temp < 250",
+        "color_temp > 454",
+        "hs_color",
+        "ct + brightness",
+        "ct + brightness (no temp support)",
+        "ct + brightness (no temp or color support)",
+        "hs + brightness",
+    ],
+)
 async def test_turn_on(
     hass: HomeAssistant,
-    mock_gateway,
+    command_store: CommandStore,
-    mock_api_factory,
+    device: Device,
-    test_features,
+    entity_id: str,
-    test_data,
+    service_data: dict[str, Any],
-    expected_result,
+    transition: dict[str, int],
-    device_id,
+    state_attributes: dict[str, Any],
 ) -> None:
     """Test turning on a light."""
-    # Note pytradfri style, not hass. Values not really important.
+    # Make sure the light is off.
-    initial_state = {
+    device.raw[ATTR_LIGHT_CONTROL][0][ATTR_DEVICE_STATE] = 0
-        "state": False,
-        "dimmer": 0,
-        "color_temp": 250,
-        "hsb_xy_color": (100, 100, 100, 100, 100),
-    }
-
-    # Setup the gateway with a mock light.
-    light = mock_light(
-        test_features=test_features, test_state=initial_state, light_number=device_id
-    )
-    mock_gateway.mock_devices.append(light)
     await setup_integration(hass)
 
-    # Use the turn_on service call to change the light state.
     await hass.services.async_call(
-        "light",
+        LIGHT_DOMAIN,
-        "turn_on",
+        SERVICE_TURN_ON,
-        {"entity_id": f"light.tradfri_light_{device_id}", **test_data},
+        {"entity_id": entity_id, **service_data, **transition},
         blocking=True,
     )
     await hass.async_block_till_done()
 
-    # Check that the light is observed.
+    await command_store.trigger_observe_callback(
-    mock_func = light.observe
+        hass, device, {ATTR_LIGHT_CONTROL: [{ATTR_DEVICE_STATE: 1}]}
-    assert len(mock_func.mock_calls) > 0
+    )
-    _, callkwargs = mock_func.call_args
-    assert "callback" in callkwargs
-    # Callback function to refresh light state.
-    callback = callkwargs["callback"]
 
-    responses = mock_gateway.mock_responses
+    state = hass.states.get(entity_id)
-    # State on command data.
+    assert state
-    data = {"3311": [{"5850": 1}]}
+    assert state.state == STATE_ON
-    # Add data for all sent commands.
+    for key, value in state_attributes.items():
-    for resp in responses:
+        # Allow some rounding error in color conversions.
-        data["3311"][0] = {**data["3311"][0], **resp["3311"][0]}
+        assert state.attributes[key] == pytest.approx(value, abs=0.01)
-
-    # Use the callback function to update the light state.
-    dev = Device(data)
-    light_data = Light(dev, 0)
-    light.light_control.lights[0] = light_data
-    callback(light)
-    await hass.async_block_till_done()
-
-    # Check that the state is correct.
-    states = hass.states.get(f"light.tradfri_light_{device_id}")
-    for result, value in expected_result.items():
-        if result == "state":
-            assert states.state == value
-        else:
-            # Allow some rounding error in color conversions.
-            assert states.attributes[result] == pytest.approx(value, abs=0.01)
 
 
-async def test_turn_off(hass: HomeAssistant, mock_gateway, mock_api_factory) -> None:
+@pytest.mark.parametrize(
+    "transition",
+    [{}, {"transition": 0}, {"transition": 1}],
+    ids=["transition_none", "transition_0", "transition_1"],
+)
+@pytest.mark.parametrize(
+    ("device", "entity_id"),
+    [
+        ("bulb_w", "light.test_w"),
+        ("bulb_ws", "light.test_ws"),
+        ("bulb_cws", "light.test_cws"),
+    ],
+    indirect=["device"],
+)
+async def test_turn_off(
+    hass: HomeAssistant,
+    command_store: CommandStore,
+    device: Device,
+    entity_id: str,
+    transition: dict[str, int],
+) -> None:
     """Test turning off a light."""
-    state = {"state": True, "dimmer": 100}
-
-    light = mock_light(test_state=state)
-    mock_gateway.mock_devices.append(light)
     await setup_integration(hass)
 
-    # Use the turn_off service call to change the light state.
     await hass.services.async_call(
-        "light", "turn_off", {"entity_id": "light.tradfri_light_0"}, blocking=True
+        LIGHT_DOMAIN,
+        SERVICE_TURN_OFF,
+        {"entity_id": entity_id, **transition},
+        blocking=True,
     )
     await hass.async_block_till_done()
 
-    # Check that the light is observed.
+    await command_store.trigger_observe_callback(
-    mock_func = light.observe
+        hass, device, {ATTR_LIGHT_CONTROL: [{ATTR_DEVICE_STATE: 0}]}
-    assert len(mock_func.mock_calls) > 0
+    )
-    _, callkwargs = mock_func.call_args
-    assert "callback" in callkwargs
-    # Callback function to refresh light state.
-    callback = callkwargs["callback"]
 
-    responses = mock_gateway.mock_responses
+    state = hass.states.get(entity_id)
-    data = {"3311": [{}]}
+    assert state
-    # Add data for all sent commands.
+    assert state.state == STATE_OFF
-    for resp in responses:
-        data["3311"][0] = {**data["3311"][0], **resp["3311"][0]}
-
-    # Use the callback function to update the light state.
-    dev = Device(data)
-    light_data = Light(dev, 0)
-    light.light_control.lights[0] = light_data
-    callback(light)
-    await hass.async_block_till_done()
-
-    # Check that the state is correct.
-    states = hass.states.get("light.tradfri_light_0")
-    assert states.state == "off"

tests/components/tradfri/test_sensor.py

@@ -1,10 +1,6 @@
 """Tradfri sensor platform tests."""
 from __future__ import annotations
 
-import json
-from typing import Any
-from unittest.mock import MagicMock, Mock
-
 import pytest
 from pytradfri.const import (
     ATTR_AIR_PURIFIER_AIR_QUALITY,
@@ -14,8 +10,6 @@ from pytradfri.const import (
     ROOT_AIR_PURIFIER,
 )
 from pytradfri.device import Device
-from pytradfri.device.air_purifier import AirPurifier
-from pytradfri.device.blind import Blind
 
 from homeassistant.components.sensor import (
     ATTR_STATE_CLASS,
@@ -37,33 +31,25 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 
 from . import GATEWAY_ID
-from .common import setup_integration, trigger_observe_callback
+from .common import CommandStore, setup_integration
 
 from tests.common import MockConfigEntry, load_fixture
 
 
 @pytest.fixture(scope="module")
-def remote_control_response() -> dict[str, Any]:
+def remote_control() -> str:
     """Return a remote control response."""
-    return json.loads(load_fixture("remote_control.json", DOMAIN))
+    return load_fixture("remote_control.json", DOMAIN)
-
-
-@pytest.fixture
-def remote_control(remote_control_response: dict[str, Any]) -> Device:
-    """Return remote control."""
-    return Device(remote_control_response)
 
 
+@pytest.mark.parametrize("device", ["remote_control"], indirect=True)
 async def test_battery_sensor(
     hass: HomeAssistant,
-    mock_gateway: Mock,
+    command_store: CommandStore,
-    mock_api_factory: MagicMock,
+    device: Device,
-    remote_control: Device,
 ) -> None:
     """Test that a battery sensor is correctly added."""
     entity_id = "sensor.test_battery"
-    device = remote_control
-    mock_gateway.mock_devices.append(device)
     await setup_integration(hass)
 
     state = hass.states.get(entity_id)
@@ -73,8 +59,8 @@ async def test_battery_sensor(
     assert state.attributes[ATTR_DEVICE_CLASS] == SensorDeviceClass.BATTERY
     assert state.attributes[ATTR_STATE_CLASS] == SensorStateClass.MEASUREMENT
 
-    await trigger_observe_callback(
+    await command_store.trigger_observe_callback(
-        hass, mock_gateway, device, {ATTR_DEVICE_INFO: {ATTR_DEVICE_BATTERY: 60}}
+        hass, device, {ATTR_DEVICE_INFO: {ATTR_DEVICE_BATTERY: 60}}
     )
 
     state = hass.states.get(entity_id)
@@ -85,16 +71,13 @@ async def test_battery_sensor(
     assert state.attributes[ATTR_STATE_CLASS] == SensorStateClass.MEASUREMENT
 
 
+@pytest.mark.parametrize("device", ["blind"], indirect=True)
 async def test_cover_battery_sensor(
     hass: HomeAssistant,
-    mock_gateway: Mock,
+    device: Device,
-    mock_api_factory: MagicMock,
-    blind: Blind,
 ) -> None:
     """Test that a battery sensor is correctly added for a cover (blind)."""
     entity_id = "sensor.test_battery"
-    device = blind.device
-    mock_gateway.mock_devices.append(device)
     await setup_integration(hass)
 
     state = hass.states.get(entity_id)
@@ -105,16 +88,14 @@ async def test_cover_battery_sensor(
     assert state.attributes[ATTR_STATE_CLASS] == SensorStateClass.MEASUREMENT
 
 
+@pytest.mark.parametrize("device", ["air_purifier"], indirect=True)
 async def test_air_quality_sensor(
     hass: HomeAssistant,
-    mock_gateway: Mock,
+    command_store: CommandStore,
-    mock_api_factory: MagicMock,
+    device: Device,
-    air_purifier: AirPurifier,
 ) -> None:
     """Test that a battery sensor is correctly added."""
     entity_id = "sensor.test_air_quality"
-    device = air_purifier.device
-    mock_gateway.mock_devices.append(device)
     await setup_integration(hass)
 
     state = hass.states.get(entity_id)
@@ -128,9 +109,8 @@ async def test_air_quality_sensor(
     assert ATTR_DEVICE_CLASS not in state.attributes
 
     # The sensor returns 65535 if the fan is turned off
-    await trigger_observe_callback(
+    await command_store.trigger_observe_callback(
         hass,
-        mock_gateway,
         device,
         {ROOT_AIR_PURIFIER: [{ATTR_AIR_PURIFIER_AIR_QUALITY: 65535}]},
     )
@@ -140,16 +120,13 @@ async def test_air_quality_sensor(
     assert state.state == STATE_UNKNOWN
 
 
+@pytest.mark.parametrize("device", ["air_purifier"], indirect=True)
 async def test_filter_time_left_sensor(
     hass: HomeAssistant,
-    mock_gateway: Mock,
+    device: Device,
-    mock_api_factory: MagicMock,
-    air_purifier: AirPurifier,
 ) -> None:
     """Test that a battery sensor is correctly added."""
     entity_id = "sensor.test_filter_time_left"
-    device = air_purifier.device
-    mock_gateway.mock_devices.append(device)
     await setup_integration(hass)
 
     state = hass.states.get(entity_id)
@@ -160,24 +137,22 @@ async def test_filter_time_left_sensor(
     assert state.attributes[ATTR_STATE_CLASS] == SensorStateClass.MEASUREMENT
 
 
+@pytest.mark.parametrize("device", ["air_purifier"], indirect=True)
 async def test_sensor_available(
     hass: HomeAssistant,
-    mock_gateway: Mock,
+    command_store: CommandStore,
-    mock_api_factory: MagicMock,
+    device: Device,
-    air_purifier: AirPurifier,
 ) -> None:
     """Test sensor available property."""
     entity_id = "sensor.test_filter_time_left"
-    device = air_purifier.device
-    mock_gateway.mock_devices.append(device)
     await setup_integration(hass)
 
     state = hass.states.get(entity_id)
     assert state
     assert state.state == "4320"
 
-    await trigger_observe_callback(
+    await command_store.trigger_observe_callback(
-        hass, mock_gateway, device, {ATTR_REACHABLE_STATE: 0}
+        hass, device, {ATTR_REACHABLE_STATE: 0}
     )
 
     state = hass.states.get(entity_id)
@@ -185,14 +160,13 @@ async def test_sensor_available(
     assert state.state == STATE_UNAVAILABLE
 
 
+@pytest.mark.parametrize("device", ["remote_control"], indirect=True)
 async def test_unique_id_migration(
     hass: HomeAssistant,
-    mock_gateway: Mock,
+    entity_registry: er.EntityRegistry,
-    mock_api_factory: MagicMock,
+    device: Device,
-    remote_control: Device,
 ) -> None:
     """Test unique ID is migrated from old format to new."""
-    ent_reg = er.async_get(hass)
     old_unique_id = f"{GATEWAY_ID}-65536"
     entry = MockConfigEntry(
         domain=DOMAIN,
@@ -209,7 +183,7 @@ async def test_unique_id_migration(
     entity_id = "sensor.test"
     entity_name = entity_id.split(".")[1]
 
-    entity_entry = ent_reg.async_get_or_create(
+    entity_entry = entity_registry.async_get_or_create(
         SENSOR_DOMAIN,
         DOMAIN,
         old_unique_id,
@@ -221,15 +195,15 @@ async def test_unique_id_migration(
     assert entity_entry.entity_id == entity_id
     assert entity_entry.unique_id == old_unique_id
 
-    # Add a sensor to the gateway so that it populates coordinator list
-    device = remote_control
-    mock_gateway.mock_devices.append(device)
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
     # Check that new RegistryEntry is using new unique ID format
     new_unique_id = f"{old_unique_id}-battery_level"
-    migrated_entity_entry = ent_reg.async_get(entity_id)
+    migrated_entity_entry = entity_registry.async_get(entity_id)
     assert migrated_entity_entry is not None
     assert migrated_entity_entry.unique_id == new_unique_id
-    assert ent_reg.async_get_entity_id(SENSOR_DOMAIN, DOMAIN, old_unique_id) is None
+    assert (
+        entity_registry.async_get_entity_id(SENSOR_DOMAIN, DOMAIN, old_unique_id)
+        is None
+    )

tests/components/tradfri/test_switch.py

@@ -1,58 +1,42 @@
 """Tradfri switch (recognised as sockets in the IKEA ecosystem) platform tests."""
 from __future__ import annotations
 
-import json
-from typing import Any
-from unittest.mock import MagicMock, Mock
-
 import pytest
 from pytradfri.const import ATTR_REACHABLE_STATE
 from pytradfri.device import Device
-from pytradfri.device.socket import Socket
 
 from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.components.tradfri.const import DOMAIN
 from homeassistant.const import STATE_OFF, STATE_ON, STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
 
-from .common import setup_integration, trigger_observe_callback
+from .common import CommandStore, setup_integration
 
 from tests.common import load_fixture
 
 
 @pytest.fixture(scope="module")
-def outlet() -> dict[str, Any]:
+def outlet() -> str:
     """Return an outlet response."""
-    return json.loads(load_fixture("outlet.json", DOMAIN))
+    return load_fixture("outlet.json", DOMAIN)
-
-
-@pytest.fixture
-def socket(outlet: dict[str, Any]) -> Socket:
-    """Return socket."""
-    device = Device(outlet)
-    socket_control = device.socket_control
-    assert socket_control
-    return socket_control.sockets[0]
 
 
+@pytest.mark.parametrize("device", ["outlet"], indirect=True)
 async def test_switch_available(
     hass: HomeAssistant,
-    mock_gateway: Mock,
+    command_store: CommandStore,
-    mock_api_factory: MagicMock,
+    device: Device,
-    socket: Socket,
 ) -> None:
     """Test switch available property."""
     entity_id = "switch.test"
-    device = socket.device
-    mock_gateway.mock_devices.append(device)
     await setup_integration(hass)
 
     state = hass.states.get(entity_id)
     assert state
     assert state.state == STATE_OFF
 
-    await trigger_observe_callback(
+    await command_store.trigger_observe_callback(
-        hass, mock_gateway, device, {ATTR_REACHABLE_STATE: 0}
+        hass, device, {ATTR_REACHABLE_STATE: 0}
     )
 
     state = hass.states.get(entity_id)
@@ -60,6 +44,7 @@ async def test_switch_available(
     assert state.state == STATE_UNAVAILABLE
 
 
+@pytest.mark.parametrize("device", ["outlet"], indirect=True)
 @pytest.mark.parametrize(
     ("service", "expected_state"),
     [
@@ -69,16 +54,13 @@ async def test_switch_available(
 )
 async def test_turn_on_off(
     hass: HomeAssistant,
-    mock_gateway: Mock,
+    command_store: CommandStore,
-    mock_api_factory: MagicMock,
+    device: Device,
-    socket: Socket,
     service: str,
     expected_state: str,
 ) -> None:
     """Test turning switch on/off."""
     entity_id = "switch.test"
-    device = socket.device
-    mock_gateway.mock_devices.append(device)
     await setup_integration(hass)
 
     state = hass.states.get(entity_id)
@@ -95,7 +77,7 @@ async def test_turn_on_off(
     )
     await hass.async_block_till_done()
 
-    await trigger_observe_callback(hass, mock_gateway, device)
+    await command_store.trigger_observe_callback(hass, device)
 
     state = hass.states.get(entity_id)
     assert state