feat: Refactor maintenance coordinator integration in button and sensor components

Author: perosb

custom_components/qvantum/__init__.py

@@ -49,7 +49,7 @@ class RuntimeData:
     """Class to hold your data."""
 
     coordinator: QvantumDataUpdateCoordinator
-    firmware_coordinator: QvantumMaintenanceCoordinator | None = None
+    maintenance_coordinator: QvantumMaintenanceCoordinator | None = None
     device: DeviceInfo | None = None
 
 

custom_components/qvantum/button.py

@@ -14,6 +14,7 @@
 
 from . import MyConfigEntry
 from .coordinator import QvantumDataUpdateCoordinator, handle_setting_update_response
+from .maintenance_coordinator import QvantumMaintenanceCoordinator
 from .entity import QvantumEntity
 
 _LOGGER = logging.getLogger(__name__)
@@ -26,12 +27,19 @@ async def async_setup_entry(
 ):
     """Set up the Button."""
     coordinator: QvantumDataUpdateCoordinator = config_entry.runtime_data.coordinator
+    maintenance_coordinator: QvantumMaintenanceCoordinator = (
+        config_entry.runtime_data.maintenance_coordinator
+    )
     device: DeviceInfo = config_entry.runtime_data.device
 
     buttons = []
     buttons.append(QvantumButtonEntity(coordinator, "extra_tap_water_60min", device))
 
-    buttons.append(QvantumButtonEntity(coordinator, "elevate_access", device))
+    buttons.append(
+        QvantumButtonEntity(
+            coordinator, "elevate_access", device, maintenance_coordinator
+        )
+    )
 
     async_add_entities(buttons)
 
@@ -46,8 +54,10 @@ def __init__(
         coordinator: QvantumDataUpdateCoordinator,
         button_key: str,
         device: DeviceInfo,
+        maintenance_coordinator: QvantumMaintenanceCoordinator | None = None,
     ) -> None:
         super().__init__(coordinator, button_key, device)
+        self._maintenance_coordinator = maintenance_coordinator
 
         if button_key == "elevate_access":
             self._attr_entity_category = EntityCategory.DIAGNOSTIC
@@ -76,6 +86,10 @@ async def async_press(self) -> None:
 
                 _LOGGER.info("Access level: %s", response)
 
+                # Refresh maintenance coordinator to update access level data
+                if self._maintenance_coordinator:
+                    await self._maintenance_coordinator.async_refresh()
+
     @property
     def available(self):
         """Check if button is available."""

custom_components/qvantum/sensor.py

@@ -78,26 +78,26 @@ async def async_setup_entry(
     )
 
     # Add firmware sensors
-    firmware_coordinator = config_entry.runtime_data.firmware_coordinator
+    maintenance_coordinator = config_entry.runtime_data.maintenance_coordinator
     sensors.append(
-        QvantumAccessExpireEntity(firmware_coordinator, "expiresAt", device, True)
+        QvantumAccessExpireEntity(maintenance_coordinator, "expiresAt", device, True)
     )
     sensors.append(
         QvantumFirmwareSensorEntity(
-            firmware_coordinator, "display_fw_version", device, True
+            maintenance_coordinator, "display_fw_version", device, True
         )
     )
     sensors.append(
-        QvantumFirmwareSensorEntity(firmware_coordinator, "cc_fw_version", device, True)
+        QvantumFirmwareSensorEntity(maintenance_coordinator, "cc_fw_version", device, True)
     )
     sensors.append(
         QvantumFirmwareSensorEntity(
-            firmware_coordinator, "inv_fw_version", device, True
+            maintenance_coordinator, "inv_fw_version", device, True
         )
     )
     sensors.append(
         QvantumFirmwareLastCheckSensorEntity(
-            firmware_coordinator, "firmware_last_check", device, True
+            maintenance_coordinator, "firmware_last_check", device, True
         )
     )
 

tests/test_button_working.py

@@ -29,6 +29,14 @@ def mock_coordinator():
     return coordinator
 
 
+@pytest.fixture
+def mock_maintenance_coordinator():
+    """Create a mock maintenance coordinator."""
+    coordinator = MagicMock()
+    coordinator.async_refresh = AsyncMock()
+    return coordinator
+
+
 @pytest.fixture
 def mock_device():
     """Create a mock device."""
@@ -57,7 +65,10 @@ def test_button_entity_initialization(self, mock_coordinator, mock_device):
 
         # Test elevate_access button
         elevate_button = QvantumButtonEntity(
-            mock_coordinator, "elevate_access", mock_device
+            mock_coordinator,
+            "elevate_access",
+            mock_device,
+            mock_maintenance_coordinator,
         )
 
         assert elevate_button._metric_key == "elevate_access"
@@ -89,43 +100,62 @@ async def test_async_press_extra_tap_water_60min(
 
     @pytest.mark.asyncio
     async def test_async_press_elevate_access(
-        self, mock_coordinator, mock_device
+        self, mock_coordinator, mock_device, mock_maintenance_coordinator
     ):
         """Test pressing the elevate access button."""
         button = QvantumButtonEntity(
-            mock_coordinator, "elevate_access", mock_device
+            mock_coordinator,
+            "elevate_access",
+            mock_device,
+            mock_maintenance_coordinator,
         )
 
         await button.async_press()
 
         mock_coordinator.api.elevate_access.assert_called_once_with(
             "test_device_123"
         )
+        # Verify maintenance coordinator is refreshed after elevating access
+        mock_maintenance_coordinator.async_refresh.assert_called_once()
 
     @pytest.mark.asyncio
     async def test_async_press_elevate_access_failure(
-        self, mock_coordinator, mock_device, caplog
+        self, mock_coordinator, mock_device, mock_maintenance_coordinator, caplog
     ):
         """Test pressing the elevate access button when elevation fails."""
         # Mock the API to return None (failure)
         mock_coordinator.api.elevate_access = AsyncMock(return_value=None)
 
-        button = QvantumButtonEntity(mock_coordinator, "elevate_access", mock_device)
+        button = QvantumButtonEntity(
+            mock_coordinator,
+            "elevate_access",
+            mock_device,
+            mock_maintenance_coordinator,
+        )
 
         await button.async_press()
 
         mock_coordinator.api.elevate_access.assert_called_once_with("test_device_123")
 
         # Verify error is logged
         assert "Failed to elevate access" in caplog.text
+        # Verify maintenance coordinator is not refreshed on failure
+        mock_maintenance_coordinator.async_refresh.assert_not_called()
 
     @pytest.mark.asyncio
     async def test_async_setup_entry(
-        self, hass, mock_config_entry, mock_coordinator, mock_device
+        self,
+        hass,
+        mock_config_entry,
+        mock_coordinator,
+        mock_device,
+        mock_maintenance_coordinator,
     ):
         """Test setting up button entities."""
         mock_config_entry.runtime_data = RuntimeData(
-            coordinator=mock_coordinator, device=mock_device
+            coordinator=mock_coordinator,
+            device=mock_device,
+            maintenance_coordinator=mock_maintenance_coordinator,
         )
 
         async_add_entities = MagicMock()