From eff82ba82c084125143c89e07bc9249d74febc72 Mon Sep 17 00:00:00 2001
From: Brett Adams <Bre77@users.noreply.github.com>
Date: Thu, 22 Feb 2024 19:50:44 +1000
Subject: [PATCH] Add wake up timeout to Teslemetry (#109037)

---
 homeassistant/components/teslemetry/entity.py | 18 ++++++-
 tests/components/teslemetry/conftest.py       | 10 ++++
 tests/components/teslemetry/test_climate.py   | 47 ++++++++++++++++++-
 3 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/homeassistant/components/teslemetry/entity.py b/homeassistant/components/teslemetry/entity.py
index d8dcf9934cc..024d0603e7e 100644
--- a/homeassistant/components/teslemetry/entity.py
+++ b/homeassistant/components/teslemetry/entity.py
@@ -3,6 +3,9 @@
 import asyncio
 from typing import Any
 
+from tesla_fleet_api.exceptions import TeslaFleetError
+
+from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
@@ -45,11 +48,22 @@ class TeslemetryVehicleEntity(CoordinatorEntity[TeslemetryVehicleDataCoordinator
     async def wake_up_if_asleep(self) -> None:
         """Wake up the vehicle if its asleep."""
         async with self._wakelock:
+            times = 0
             while self.coordinator.data["state"] != TeslemetryState.ONLINE:
-                state = (await self.api.wake_up())["response"]["state"]
+                try:
+                    if times == 0:
+                        cmd = await self.api.wake_up()
+                    else:
+                        cmd = await self.api.vehicle()
+                    state = cmd["response"]["state"]
+                except TeslaFleetError as e:
+                    raise HomeAssistantError(str(e)) from e
                 self.coordinator.data["state"] = state
                 if state != TeslemetryState.ONLINE:
-                    await asyncio.sleep(5)
+                    times += 1
+                    if times >= 4:  # Give up after 30 seconds total
+                        raise HomeAssistantError("Could not wake up vehicle")
+                    await asyncio.sleep(times * 5)
 
     def get(self, key: str | None = None, default: Any | None = None) -> Any:
         """Return a specific value from coordinator data."""
diff --git a/tests/components/teslemetry/conftest.py b/tests/components/teslemetry/conftest.py
index 0fc279eaa21..8c1fe070dde 100644
--- a/tests/components/teslemetry/conftest.py
+++ b/tests/components/teslemetry/conftest.py
@@ -37,6 +37,16 @@ def mock_wake_up():
         yield mock_wake_up
 
 
+@pytest.fixture(autouse=True)
+def mock_vehicle():
+    """Mock Tesla Fleet API Vehicle Specific vehicle method."""
+    with patch(
+        "homeassistant.components.teslemetry.VehicleSpecific.vehicle",
+        return_value=WAKE_UP_ONLINE,
+    ) as mock_vehicle:
+        yield mock_vehicle
+
+
 @pytest.fixture(autouse=True)
 def mock_request():
     """Mock Tesla Fleet API Vehicle Specific class."""
diff --git a/tests/components/teslemetry/test_climate.py b/tests/components/teslemetry/test_climate.py
index ede38a695e2..2e791f68b93 100644
--- a/tests/components/teslemetry/test_climate.py
+++ b/tests/components/teslemetry/test_climate.py
@@ -26,6 +26,7 @@ from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import entity_registry as er
 
 from . import assert_entities, setup_platform
+from .const import WAKE_UP_ASLEEP, WAKE_UP_ONLINE
 
 from tests.common import async_fire_time_changed
 
@@ -108,7 +109,11 @@ async def test_errors(
 
 
 async def test_asleep_or_offline(
-    hass: HomeAssistant, mock_vehicle_data, freezer: FrozenDateTimeFactory
+    hass: HomeAssistant,
+    mock_vehicle_data,
+    mock_wake_up,
+    mock_vehicle,
+    freezer: FrozenDateTimeFactory,
 ) -> None:
     """Tests asleep is handled."""
 
@@ -123,9 +128,47 @@ async def test_asleep_or_offline(
     async_fire_time_changed(hass)
     await hass.async_block_till_done()
     mock_vehicle_data.assert_called_once()
+    mock_wake_up.reset_mock()
 
-    # Run a command that will wake up the vehicle, but not immediately
+    # Run a command but fail trying to wake up the vehicle
+    mock_wake_up.side_effect = InvalidCommand
+    with pytest.raises(HomeAssistantError) as error:
+        await hass.services.async_call(
+            CLIMATE_DOMAIN,
+            SERVICE_TURN_ON,
+            {ATTR_ENTITY_ID: [entity_id]},
+            blocking=True,
+        )
+        assert error
+    mock_wake_up.assert_called_once()
+
+    mock_wake_up.side_effect = None
+    mock_wake_up.reset_mock()
+
+    # Run a command but timeout trying to wake up the vehicle
+    mock_wake_up.return_value = WAKE_UP_ASLEEP
+    mock_vehicle.return_value = WAKE_UP_ASLEEP
+    with patch(
+        "homeassistant.components.teslemetry.entity.asyncio.sleep"
+    ), pytest.raises(HomeAssistantError) as error:
+        await hass.services.async_call(
+            CLIMATE_DOMAIN,
+            SERVICE_TURN_ON,
+            {ATTR_ENTITY_ID: [entity_id]},
+            blocking=True,
+        )
+        assert error
+    mock_wake_up.assert_called_once()
+    mock_vehicle.assert_called()
+
+    mock_wake_up.reset_mock()
+    mock_vehicle.reset_mock()
+    mock_wake_up.return_value = WAKE_UP_ONLINE
+    mock_vehicle.return_value = WAKE_UP_ONLINE
+
+    # Run a command and wake up the vehicle immediately
     await hass.services.async_call(
         CLIMATE_DOMAIN, SERVICE_TURN_ON, {ATTR_ENTITY_ID: [entity_id]}, blocking=True
     )
     await hass.async_block_till_done()
+    mock_wake_up.assert_called_once()