Ensure togrill detects disconnected devices (home-assistant#153067)

Author: elupus

homeassistant/components/togrill/coordinator.py

@@ -139,7 +139,11 @@ async def _connect_and_update_registry(self) -> Client:
             raise DeviceNotFound("Unable to find device")
 
         try:
-            client = await Client.connect(device, self._notify_callback)
+            client = await Client.connect(
+                device,
+                self._notify_callback,
+                disconnected_callback=self._disconnected_callback,
+            )
         except BleakError as exc:
             self.logger.debug("Connection failed", exc_info=True)
             raise DeviceNotFound("Unable to connect to device") from exc
@@ -169,9 +173,6 @@ async def async_shutdown(self) -> None:
         self.client = None
 
     async def _get_connected_client(self) -> Client:
-        if self.client and not self.client.is_connected:
-            await self.client.disconnect()
-            self.client = None
         if self.client:
             return self.client
 
@@ -196,6 +197,12 @@ def _notify_callback(self, packet: Packet):
 
     async def _async_update_data(self) -> dict[tuple[int, int | None], Packet]:
         """Poll the device."""
+        if self.client and not self.client.is_connected:
+            await self.client.disconnect()
+            self.client = None
+            self._async_request_refresh_soon()
+            raise DeviceFailed("Device was disconnected")
+
         client = await self._get_connected_client()
         try:
             await client.request(PacketA0Notify)
@@ -206,12 +213,23 @@ async def _async_update_data(self) -> dict[tuple[int, int | None], Packet]:
             raise DeviceFailed(f"Device failed {exc}") from exc
         return self.data
 
+    @callback
+    def _async_request_refresh_soon(self) -> None:
+        self.config_entry.async_create_task(
+            self.hass, self.async_request_refresh(), eager_start=False
+        )
+
+    @callback
+    def _disconnected_callback(self) -> None:
+        """Handle Bluetooth device being disconnected."""
+        self._async_request_refresh_soon()
+
     @callback
     def _async_handle_bluetooth_event(
         self,
         service_info: BluetoothServiceInfoBleak,
         change: BluetoothChange,
     ) -> None:
         """Handle a Bluetooth event."""
-        if not self.client and isinstance(self.last_exception, DeviceNotFound):
-            self.hass.async_create_task(self.async_refresh())
+        if isinstance(self.last_exception, DeviceNotFound):
+            self._async_request_refresh_soon()

tests/components/togrill/conftest.py

@@ -57,9 +57,18 @@ def mock_client(enable_bluetooth: None, mock_client_class: Mock) -> Generator[Mo
     client_object.mocked_notify = None
 
     async def _connect(
-        address: str, callback: Callable[[Packet], None] | None = None
+        address: str,
+        callback: Callable[[Packet], None] | None = None,
+        disconnected_callback: Callable[[], None] | None = None,
     ) -> Mock:
         client_object.mocked_notify = callback
+        if disconnected_callback:
+
+            def _disconnected_callback():
+                client_object.is_connected = False
+                disconnected_callback()
+
+            client_object.mocked_disconnected_callback = _disconnected_callback
         return client_object
 
     async def _disconnect() -> None:

tests/components/togrill/test_sensor.py

@@ -1,7 +1,8 @@
 """Test sensors for ToGrill integration."""
 
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
+from habluetooth import BluetoothServiceInfoBleak
 import pytest
 from syrupy.assertion import SnapshotAssertion
 from togrill_bluetooth.packets import PacketA0Notify, PacketA1Notify
@@ -16,6 +17,16 @@
 from tests.components.bluetooth import inject_bluetooth_service_info
 
 
+def patch_async_ble_device_from_address(
+    return_value: BluetoothServiceInfoBleak | None = None,
+):
+    """Patch async_ble_device_from_address to return a mocked BluetoothServiceInfoBleak."""
+    return patch(
+        "homeassistant.components.bluetooth.async_ble_device_from_address",
+        return_value=return_value,
+    )
+
+
 @pytest.mark.parametrize(
     "packets",
     [
@@ -57,3 +68,51 @@ async def test_setup(
         mock_client.mocked_notify(packet)
 
     await snapshot_platform(hass, entity_registry, snapshot, mock_entry.entry_id)
+
+
+async def test_device_disconnected(
+    hass: HomeAssistant,
+    mock_entry: MockConfigEntry,
+    mock_client: Mock,
+) -> None:
+    """Test the switch set."""
+    inject_bluetooth_service_info(hass, TOGRILL_SERVICE_INFO)
+
+    await setup_entry(hass, mock_entry, [Platform.SENSOR])
+
+    entity_id = "sensor.pro_05_battery"
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == "0"
+
+    with patch_async_ble_device_from_address():
+        mock_client.mocked_disconnected_callback()
+        await hass.async_block_till_done()
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == "unavailable"
+
+
+async def test_device_discovered(
+    hass: HomeAssistant,
+    mock_entry: MockConfigEntry,
+    mock_client: Mock,
+) -> None:
+    """Test the switch set."""
+
+    await setup_entry(hass, mock_entry, [Platform.SENSOR])
+
+    entity_id = "sensor.pro_05_battery"
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == "unavailable"
+
+    inject_bluetooth_service_info(hass, TOGRILL_SERVICE_INFO)
+    await hass.async_block_till_done()
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == "0"