From da9fc495ca9f5b526f53e94de78f069d7342d360 Mon Sep 17 00:00:00 2001
From: Joost Lekkerkerker <joostlek@outlook.com>
Date: Fri, 25 Aug 2023 11:19:40 +0200
Subject: [PATCH] Improve SRP Energy coordinator (#99010)

* Improve SRP Energy coordinator

* Use time instead of asyncio
---
 .../components/srp_energy/coordinator.py      | 54 +++++++++----------
 tests/components/srp_energy/test_sensor.py    | 29 ++++++++--
 2 files changed, 51 insertions(+), 32 deletions(-)

diff --git a/homeassistant/components/srp_energy/coordinator.py b/homeassistant/components/srp_energy/coordinator.py
index 9637f176886..a72ea4d3334 100644
--- a/homeassistant/components/srp_energy/coordinator.py
+++ b/homeassistant/components/srp_energy/coordinator.py
@@ -4,7 +4,6 @@ from __future__ import annotations
 import asyncio
 from datetime import timedelta
 
-from requests.exceptions import ConnectionError as ConnectError, HTTPError, Timeout
 from srpenergy.client import SrpEnergyClient
 
 from homeassistant.config_entries import ConfigEntry
@@ -14,6 +13,8 @@ from homeassistant.util import dt as dt_util
 
 from .const import DOMAIN, LOGGER, MIN_TIME_BETWEEN_UPDATES, PHOENIX_TIME_ZONE
 
+TIMEOUT = 10
+
 
 class SRPEnergyDataUpdateCoordinator(DataUpdateCoordinator[float]):
     """A srp_energy Data Update Coordinator."""
@@ -40,38 +41,35 @@ class SRPEnergyDataUpdateCoordinator(DataUpdateCoordinator[float]):
         so entities can quickly look up their data.
         """
         LOGGER.debug("async_update_data enter")
+        # Fetch srp_energy data
+        phx_time_zone = dt_util.get_time_zone(PHOENIX_TIME_ZONE)
+        end_date = dt_util.now(phx_time_zone)
+        start_date = end_date - timedelta(days=1)
         try:
-            # Fetch srp_energy data
-            phx_time_zone = dt_util.get_time_zone(PHOENIX_TIME_ZONE)
-            end_date = dt_util.now(phx_time_zone)
-            start_date = end_date - timedelta(days=1)
-
-            async with asyncio.timeout(10):
+            async with asyncio.timeout(TIMEOUT):
                 hourly_usage = await self.hass.async_add_executor_job(
                     self._client.usage,
                     start_date,
                     end_date,
                     self._is_time_of_use,
                 )
-
-                LOGGER.debug(
-                    "async_update_data: Received %s record(s) from %s to %s",
-                    len(hourly_usage) if hourly_usage else "None",
-                    start_date,
-                    end_date,
-                )
-
-                previous_daily_usage = 0.0
-                for _, _, _, kwh, _ in hourly_usage:
-                    previous_daily_usage += float(kwh)
-
-                LOGGER.debug(
-                    "async_update_data: previous_daily_usage %s",
-                    previous_daily_usage,
-                )
-
-                return previous_daily_usage
-        except TimeoutError as timeout_err:
-            raise UpdateFailed("Timeout communicating with API") from timeout_err
-        except (ConnectError, HTTPError, Timeout, ValueError, TypeError) as err:
+        except (ValueError, TypeError) as err:
             raise UpdateFailed(f"Error communicating with API: {err}") from err
+
+        LOGGER.debug(
+            "async_update_data: Received %s record(s) from %s to %s",
+            len(hourly_usage) if hourly_usage else "None",
+            start_date,
+            end_date,
+        )
+
+        previous_daily_usage = 0.0
+        for _, _, _, kwh, _ in hourly_usage:
+            previous_daily_usage += float(kwh)
+
+        LOGGER.debug(
+            "async_update_data: previous_daily_usage %s",
+            previous_daily_usage,
+        )
+
+        return previous_daily_usage
diff --git a/tests/components/srp_energy/test_sensor.py b/tests/components/srp_energy/test_sensor.py
index 0e3075c6ac8..1ae213e4bf1 100644
--- a/tests/components/srp_energy/test_sensor.py
+++ b/tests/components/srp_energy/test_sensor.py
@@ -1,7 +1,7 @@
 """Tests for the srp_energy sensor platform."""
+import time
 from unittest.mock import patch
 
-import pytest
 from requests.models import HTTPError
 
 from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass
@@ -46,10 +46,8 @@ async def test_srp_entity(hass: HomeAssistant, init_integration) -> None:
     assert usage_state.attributes.get(ATTR_ICON) == "mdi:flash"
 
 
-@pytest.mark.parametrize("error", [TimeoutError, HTTPError])
 async def test_srp_entity_update_failed(
     hass: HomeAssistant,
-    error: Exception,
     mock_config_entry: MockConfigEntry,
 ) -> None:
     """Test the SrpEntity."""
@@ -59,7 +57,30 @@ async def test_srp_entity_update_failed(
     ) as srp_energy_mock:
         client = srp_energy_mock.return_value
         client.validate.return_value = True
-        client.usage.side_effect = error
+        client.usage.side_effect = HTTPError
+        mock_config_entry.add_to_hass(hass)
+
+        await hass.config_entries.async_setup(mock_config_entry.entry_id)
+        await hass.async_block_till_done()
+
+    usage_state = hass.states.get("sensor.home_energy_usage")
+    assert usage_state is None
+
+
+async def test_srp_entity_timeout(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test the SrpEntity timing out."""
+
+    with patch(
+        "homeassistant.components.srp_energy.SrpEnergyClient", autospec=True
+    ) as srp_energy_mock, patch(
+        "homeassistant.components.srp_energy.coordinator.TIMEOUT", 0
+    ):
+        client = srp_energy_mock.return_value
+        client.validate.return_value = True
+        client.usage = lambda _, __, ___: time.sleep(1)
         mock_config_entry.add_to_hass(hass)
 
         await hass.config_entries.async_setup(mock_config_entry.entry_id)