Fix capabilities detection.

denpamusic · Aug 4, 2022 · 1ab3417 · 1ab3417
1 parent 88cf9e7
commit 1ab3417
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/custom_components/plum_ecomax/connection.py b/custom_components/plum_ecomax/connection.py
@@ -74,11 +74,10 @@ async def async_check_connection(
 @timeout(seconds=DEFAULT_TIMEOUT)
 async def async_get_device_capabilities(device: Device) -> list[str]:
     """Return device capabilities, presented as list of allowed keys."""
-    sensors = await device.get_value("sensors")
-    parameters = await device.get_value("parameters")
+    await device.get_value("sensors")
+    await device.get_value("parameters")
     capabilities = ["product", "modules"]
-    capabilities += list(sensors.keys())
-    capabilities += list(parameters.keys())
+    capabilities += list(device.data.keys())
     for capability in ("fuel_burned", "boiler_control", "password"):
         try:
             await device.get_value(capability, timeout=5)

diff --git a/tests/test_connection.py b/tests/test_connection.py
@@ -54,12 +54,17 @@ async def test_async_check_connection() -> None:
     mock_device.get_value.side_effect = (
         mock_product,
         mock_modules,
-        {"test_sensor": "test_value", "water_heater_temp": 50},
-        {"test_parameter": "test_value"},
+        True,
+        True,
         "fuel_burned",
         "boiler_control",
         asyncio.TimeoutError,
     )
+    mock_device.data = {
+        "test_sensor": "test_value",
+        "water_heater_temp": 50,
+        "test_parameter": "test_value",
+    }
     result = await async_check_connection(mock_connection)
     calls = (
         call("product"),