Fix some tests but not all

tests/components/tradfri/common.py

@@ -5,8 +5,7 @@ from dataclasses import dataclass
 from typing import Any
 
 from pytradfri.command import Command
-from pytradfri.const import ATTR_ID
+from pytradfri.device import Device, DeviceResponse
-from pytradfri.device import Device
 from pytradfri.gateway import Gateway
 
 from homeassistant.components import tradfri
@@ -25,13 +24,13 @@ class CommandStore:
     mock_responses: dict[str, Any]
 
     def register_device(
-        self, gateway: Gateway, device_response: dict[str, Any]
+        self, gateway: Gateway, device_response: DeviceResponse
     ) -> None:
         """Register device response."""
         get_devices_command = gateway.get_devices()
-        self.register_response(get_devices_command, [device_response[ATTR_ID]])
+        self.register_response(get_devices_command, [device_response.id])
-        get_device_command = gateway.get_device(device_response[ATTR_ID])
+        get_device_command = gateway.get_device(device_response.id)
-        self.register_response(get_device_command, device_response)
+        self.register_response(get_device_command, device_response.dict())
 
     def register_response(self, command: Command, response: Any) -> None:
         """Register command response."""

tests/components/tradfri/conftest.py

@@ -9,13 +9,12 @@ 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.gateway import Gateway
 
 from homeassistant.components.tradfri.const import DOMAIN
 
-from . import GATEWAY_ID, TRADFRI_PATH
+from . import TRADFRI_PATH
 from .common import CommandStore
 
 from tests.common import load_fixture
@@ -30,12 +29,12 @@ def mock_entry_setup() -> Generator[AsyncMock]:
 
 
 @pytest.fixture(name="mock_gateway", autouse=True)
-def mock_gateway_fixture(command_store: CommandStore) -> Gateway:
+def mock_gateway_fixture(command_store: CommandStore, gateway_response: str) -> Gateway:
     """Mock a Tradfri gateway."""
     gateway = Gateway()
     command_store.register_response(
         gateway.get_gateway_info(),
-        {ATTR_GATEWAY_ID: GATEWAY_ID, ATTR_FIRMWARE_VERSION: "1.2.1234"},
+        json.loads(gateway_response),
     )
     command_store.register_response(
         gateway.get_devices(),
@@ -96,6 +95,12 @@ def device(
     return device
 
 
+@pytest.fixture(scope="package")
+def gateway_response() -> str:
+    """Return a gateway response."""
+    return load_fixture("gateway.json", DOMAIN)
+
+
 @pytest.fixture(scope="package")
 def air_purifier() -> str:
     """Return an air purifier response."""

tests/components/tradfri/fixtures/gateway.json

@@ -0,0 +1,29 @@
+{
+  "9023": "xyz.pool.ntp.pool",
+  "9059": 1509788799,
+  "9073": 0,
+  "9060": "2017-11-04T09:46:39.046784Z",
+  "9080": 0,
+  "9071": 1,
+  "9062": 0,
+  "9061": 0,
+  "9093": 0,
+  "9029": "1.2.1234",
+  "9081": "mock-gateway-id",
+  "9092": 0,
+  "9069": 1509474847,
+  "9082": true,
+  "9055": 0,
+  "9083": "123-45-67",
+  "9066": 5,
+  "9054": 0,
+  "9077": 0,
+  "9072": 0,
+  "9074": 0,
+  "9075": 0,
+  "9076": 0,
+  "9078": 0,
+  "9079": 0,
+  "9106": 0,
+  "9105": 0
+}

tests/components/tradfri/test_light.py

@@ -246,7 +246,7 @@ async def test_turn_on(
 ) -> None:
     """Test turning on a light."""
     # Make sure the light is off.
-    device.raw[ATTR_LIGHT_CONTROL][0][ATTR_DEVICE_STATE] = 0
+    device.raw.light_control[0].state = 0
     await setup_integration(hass)
 
     await hass.services.async_call(