From a33aacfcaa109257ea5c83d39531a1efa9485e4a Mon Sep 17 00:00:00 2001
From: Nalin Mahajan <nalinmahajan@outlook.com>
Date: Tue, 23 Apr 2024 15:10:16 -0500
Subject: [PATCH] Add Retry for C4 API due to flakiness (#113857)

Co-authored-by: nalin29 <nalin29@github.com>
---
 homeassistant/components/control4/__init__.py | 26 ++++++++++++++-----
 homeassistant/components/control4/const.py    |  2 ++
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/control4/__init__.py b/homeassistant/components/control4/__init__.py
index b8d195fcb05..4b24ac6bf77 100644
--- a/homeassistant/components/control4/__init__.py
+++ b/homeassistant/components/control4/__init__.py
@@ -30,6 +30,7 @@ from homeassistant.helpers.update_coordinator import (
 )
 
 from .const import (
+    API_RETRY_TIMES,
     CONF_ACCOUNT,
     CONF_CONFIG_LISTENER,
     CONF_CONTROLLER_UNIQUE_ID,
@@ -47,6 +48,18 @@ _LOGGER = logging.getLogger(__name__)
 PLATFORMS = [Platform.LIGHT, Platform.MEDIA_PLAYER]
 
 
+async def call_c4_api_retry(func, *func_args):
+    """Call C4 API function and retry on failure."""
+    for i in range(API_RETRY_TIMES):
+        try:
+            output = await func(*func_args)
+            return output
+        except client_exceptions.ClientError as exception:
+            _LOGGER.error("Error connecting to Control4 account API: %s", exception)
+            if i == API_RETRY_TIMES - 1:
+                raise ConfigEntryNotReady(exception) from exception
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Control4 from a config entry."""
     hass.data.setdefault(DOMAIN, {})
@@ -74,18 +87,19 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     controller_unique_id = config[CONF_CONTROLLER_UNIQUE_ID]
     entry_data[CONF_CONTROLLER_UNIQUE_ID] = controller_unique_id
 
-    director_token_dict = await account.getDirectorBearerToken(controller_unique_id)
-    director_session = aiohttp_client.async_get_clientsession(hass, verify_ssl=False)
+    director_token_dict = await call_c4_api_retry(
+        account.getDirectorBearerToken, controller_unique_id
+    )
 
+    director_session = aiohttp_client.async_get_clientsession(hass, verify_ssl=False)
     director = C4Director(
         config[CONF_HOST], director_token_dict[CONF_TOKEN], director_session
     )
     entry_data[CONF_DIRECTOR] = director
 
-    # Add Control4 controller to device registry
-    controller_href = (await account.getAccountControllers())["href"]
-    entry_data[CONF_DIRECTOR_SW_VERSION] = await account.getControllerOSVersion(
-        controller_href
+    controller_href = (await call_c4_api_retry(account.getAccountControllers))["href"]
+    entry_data[CONF_DIRECTOR_SW_VERSION] = await call_c4_api_retry(
+        account.getControllerOSVersion, controller_href
     )
 
     _, model, mac_address = controller_unique_id.split("_", 3)
diff --git a/homeassistant/components/control4/const.py b/homeassistant/components/control4/const.py
index f8d939e1ac5..57074c00108 100644
--- a/homeassistant/components/control4/const.py
+++ b/homeassistant/components/control4/const.py
@@ -5,6 +5,8 @@ DOMAIN = "control4"
 DEFAULT_SCAN_INTERVAL = 5
 MIN_SCAN_INTERVAL = 1
 
+API_RETRY_TIMES = 5
+
 CONF_ACCOUNT = "account"
 CONF_DIRECTOR = "director"
 CONF_DIRECTOR_SW_VERSION = "director_sw_version"