From b6e76c25a3ac60fbfdd7f72f4c9c5ba9b8319649 Mon Sep 17 00:00:00 2001
From: Andrew Jackson <andrew@codechimp.org>
Date: Wed, 20 Nov 2024 16:46:36 +0000
Subject: [PATCH] Custom library url (#2416)

* Add custom library url

* Update library url
---
 custom_components/battery_notes/__init__.py   |  7 +++++++
 custom_components/battery_notes/const.py      |  3 ++-
 .../battery_notes/library_updater.py          | 20 ++++++++++++++-----
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/custom_components/battery_notes/__init__.py b/custom_components/battery_notes/__init__.py
index 6925d67d7..4329d184b 100644
--- a/custom_components/battery_notes/__init__.py
+++ b/custom_components/battery_notes/__init__.py
@@ -31,6 +31,7 @@
     CONF_ENABLE_AUTODISCOVERY,
     CONF_ENABLE_REPLACED,
     CONF_HIDE_BATTERY,
+    CONF_LIBRARY_URL,
     CONF_ROUND_BATTERY,
     CONF_SHOW_ALL_DEVICES,
     CONF_USER_LIBRARY,
@@ -39,6 +40,7 @@
     DATA_STORE,
     DEFAULT_BATTERY_INCREASE_THRESHOLD,
     DEFAULT_BATTERY_LOW_THRESHOLD,
+    DEFAULT_LIBRARY_URL,
     DOMAIN,
     DOMAIN_CONFIG,
     MIN_HA_VERSION,
@@ -75,6 +77,10 @@
                         CONF_BATTERY_INCREASE_THRESHOLD,
                         default=DEFAULT_BATTERY_INCREASE_THRESHOLD,
                     ): cv.positive_int,
+                    vol.Optional(
+                        CONF_LIBRARY_URL,
+                        default=DEFAULT_LIBRARY_URL,
+                    ): cv.string,
                 },
             ),
         ),
@@ -111,6 +117,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
         CONF_ROUND_BATTERY: False,
         CONF_DEFAULT_BATTERY_LOW_THRESHOLD: DEFAULT_BATTERY_LOW_THRESHOLD,
         CONF_BATTERY_INCREASE_THRESHOLD: DEFAULT_BATTERY_INCREASE_THRESHOLD,
+        CONF_LIBRARY_URL: DEFAULT_LIBRARY_URL,
     }
 
     hass.data[DOMAIN] = {
diff --git a/custom_components/battery_notes/const.py b/custom_components/battery_notes/const.py
index 9a3e608ff..b8123096a 100644
--- a/custom_components/battery_notes/const.py
+++ b/custom_components/battery_notes/const.py
@@ -30,6 +30,7 @@
 
 DEFAULT_BATTERY_LOW_THRESHOLD = 10
 DEFAULT_BATTERY_INCREASE_THRESHOLD = 25
+DEFAULT_LIBRARY_URL = "https://raw.githubusercontent.com/andrew-codechimp/HA-Battery-Notes/main/custom_components/battery_notes/data/library.json"  # pylint: disable=line-too-long
 
 CONF_SOURCE_ENTITY_ID = "source_entity_id"
 CONF_BATTERY_TYPE = "battery_type"
@@ -42,7 +43,7 @@
 CONF_MODEL_ID = "model_id"
 CONF_MANUFACTURER = "manufacturer"
 CONF_DEVICE_NAME = "device_name"
-CONF_LIBRARY_URL = "https://raw.githubusercontent.com/andrew-codechimp/HA-Battery-Notes/main/custom_components/battery_notes/data/library.json"  # pylint: disable=line-too-long
+CONF_LIBRARY_URL = "library_url"
 CONF_SHOW_ALL_DEVICES = "show_all_devices"
 CONF_ENABLE_REPLACED = "enable_replaced"
 CONF_DEFAULT_BATTERY_LOW_THRESHOLD = "default_battery_low_threshold"
diff --git a/custom_components/battery_notes/library_updater.py b/custom_components/battery_notes/library_updater.py
index cb3f772a5..5c03a8bc3 100644
--- a/custom_components/battery_notes/library_updater.py
+++ b/custom_components/battery_notes/library_updater.py
@@ -20,6 +20,7 @@
     CONF_ENABLE_AUTODISCOVERY,
     CONF_LIBRARY_URL,
     DATA_LIBRARY_LAST_UPDATE,
+    DEFAULT_LIBRARY_URL,
     DOMAIN,
     DOMAIN_CONFIG,
 )
@@ -44,7 +45,14 @@ class LibraryUpdater:
     def __init__(self, hass: HomeAssistant):
         """Initialize the library updater."""
         self.hass = hass
-        self._client = LibraryUpdaterClient(session=async_get_clientsession(hass))
+
+        if DOMAIN_CONFIG in self.hass.data[DOMAIN]:
+            domain_config: dict = self.hass.data[DOMAIN][DOMAIN_CONFIG]
+            url = domain_config.get(CONF_LIBRARY_URL, DEFAULT_LIBRARY_URL)
+        else:
+            url = DEFAULT_LIBRARY_URL
+
+        self._client = LibraryUpdaterClient(library_url=url, session=async_get_clientsession(hass))
 
         # Fire the library check every 24 hours from now
         async_track_utc_time_change(
@@ -103,8 +111,7 @@ def _update_library_json(library_file: str, content: str) -> dict[str, Any]:
 
         except LibraryUpdaterClientError:
             _LOGGER.warning(
-                "Unable to update library, this could be a GitHub or internet "
-                "connectivity issue, will retry later."
+                "Unable to update library, will retry later."
             )
 
     async def time_to_update_library(self) -> bool:
@@ -146,14 +153,17 @@ class LibraryUpdaterClient:
 
     def __init__(
         self,
+        library_url: str,
         session: aiohttp.ClientSession,
     ) -> None:
         """Client to get latest library file from GitHub."""
+        self._library_url = library_url
         self._session = session
 
     async def async_get_data(self) -> any:
-        """Get data from the API."""
-        return await self._api_wrapper(method="get", url=CONF_LIBRARY_URL)
+        """Get data from the hosted library."""
+        _LOGGER.debug(f"Updating library from {self._library_url}")
+        return await self._api_wrapper(method="get", url=self._library_url)
 
     async def _api_wrapper(
         self,