From eb815994003ce54dd17c71c855c24d7841dd930d Mon Sep 17 00:00:00 2001
From: Scott K Logan <logans@cottsay.net>
Date: Tue, 12 Mar 2024 11:43:25 -0500
Subject: [PATCH] Fix some handle leaks in rainforest_raven (#113035)

There were leaks when
* The component was shutdown
* There was a timeout during the initial device opening

Additionally, the device was not closed/reopened when there was a
timeout reading regular data.
---
 .../rainforest_raven/coordinator.py           | 30 ++++++++++++-------
 .../rainforest_raven/test_coordinator.py      | 16 ++++++++++
 2 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/homeassistant/components/rainforest_raven/coordinator.py b/homeassistant/components/rainforest_raven/coordinator.py
index 7f2d3399a47..37e44b12eba 100644
--- a/homeassistant/components/rainforest_raven/coordinator.py
+++ b/homeassistant/components/rainforest_raven/coordinator.py
@@ -133,16 +133,27 @@ class RAVEnDataCoordinator(DataUpdateCoordinator):
             )
         return None
 
+    async def async_shutdown(self) -> None:
+        """Shutdown the coordinator."""
+        await self._cleanup_device()
+        await super().async_shutdown()
+
     async def _async_update_data(self) -> dict[str, Any]:
         try:
             device = await self._get_device()
             async with asyncio.timeout(5):
                 return await _get_all_data(device, self.entry.data[CONF_MAC])
         except RAVEnConnectionError as err:
-            if self._raven_device:
-                await self._raven_device.close()
-                self._raven_device = None
+            await self._cleanup_device()
             raise UpdateFailed(f"RAVEnConnectionError: {err}") from err
+        except TimeoutError:
+            await self._cleanup_device()
+            raise
+
+    async def _cleanup_device(self) -> None:
+        device, self._raven_device = self._raven_device, None
+        if device is not None:
+            await device.close()
 
     async def _get_device(self) -> RAVEnSerialDevice:
         if self._raven_device is not None:
@@ -150,15 +161,14 @@ class RAVEnDataCoordinator(DataUpdateCoordinator):
 
         device = RAVEnSerialDevice(self.entry.data[CONF_DEVICE])
 
-        async with asyncio.timeout(5):
-            await device.open()
-
-            try:
+        try:
+            async with asyncio.timeout(5):
+                await device.open()
                 await device.synchronize()
                 self._device_info = await device.get_device_info()
-            except Exception:
-                await device.close()
-                raise
+        except:
+            await device.close()
+            raise
 
         self._raven_device = device
         return device
diff --git a/tests/components/rainforest_raven/test_coordinator.py b/tests/components/rainforest_raven/test_coordinator.py
index 0c716aef6fc..1a5f4d3d3f7 100644
--- a/tests/components/rainforest_raven/test_coordinator.py
+++ b/tests/components/rainforest_raven/test_coordinator.py
@@ -1,5 +1,8 @@
 """Tests for the Rainforest RAVEn data coordinator."""
 
+import asyncio
+import functools
+
 from aioraven.device import RAVEnConnectionError
 import pytest
 
@@ -84,6 +87,19 @@ async def test_coordinator_device_error_update(hass: HomeAssistant, mock_device)
     assert coordinator.last_update_success is False
 
 
+async def test_coordinator_device_timeout_update(hass: HomeAssistant, mock_device):
+    """Test handling of a device timeout during an update."""
+    entry = create_mock_entry()
+    coordinator = RAVEnDataCoordinator(hass, entry)
+
+    await coordinator.async_config_entry_first_refresh()
+    assert coordinator.last_update_success is True
+
+    mock_device.get_network_info.side_effect = functools.partial(asyncio.sleep, 10)
+    await coordinator.async_refresh()
+    assert coordinator.last_update_success is False
+
+
 async def test_coordinator_comm_error(hass: HomeAssistant, mock_device):
     """Test handling of an error parsing or reading raw device data."""
     entry = create_mock_entry()