From a9f1537e2167e959a05f5db0def85b9666f28858 Mon Sep 17 00:00:00 2001
From: Alexei Chetroi <lexoid@gmail.com>
Date: Wed, 30 Sep 2020 16:52:04 -0400
Subject: [PATCH] Allow configuration of IKEA update url (#508)

---
 zigpy/config/__init__.py | 2 ++
 zigpy/ota/provider.py    | 7 +++++--
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/zigpy/config/__init__.py b/zigpy/config/__init__.py
index 64bf3ada1..74fe3a2eb 100644
--- a/zigpy/config/__init__.py
+++ b/zigpy/config/__init__.py
@@ -36,6 +36,7 @@
 CONF_OTA = "ota"
 CONF_OTA_DIR = "otau_directory"
 CONF_OTA_IKEA = "ikea_provider"
+CONF_OTA_IKEA_URL = "ikea_update_url"
 CONF_OTA_LEDVANCE = "ledvance_provider"
 CONF_TOPO_SCAN_PERIOD = "topology_scan_period"
 
@@ -72,6 +73,7 @@
 SCHEMA_OTA = {
     vol.Optional(CONF_OTA_DIR, default=CONF_OTA_OTAU_DIR_DEFAULT): vol.Any(None, str),
     vol.Optional(CONF_OTA_IKEA, default=CONF_OTA_IKEA_DEFAULT): cv_boolean,
+    vol.Optional(CONF_OTA_IKEA_URL): vol.Url(),
     vol.Optional(CONF_OTA_LEDVANCE, default=CONF_OTA_LEDVANCE_DEFAULT): cv_boolean,
 }
 
diff --git a/zigpy/ota/provider.py b/zigpy/ota/provider.py
index 7f8e084fb..3b53452ca 100644
--- a/zigpy/ota/provider.py
+++ b/zigpy/ota/provider.py
@@ -11,7 +11,7 @@
 import aiohttp
 import attr
 
-from zigpy.config import CONF_OTA_DIR
+from zigpy.config import CONF_OTA_DIR, CONF_OTA_IKEA_URL
 from zigpy.ota.image import ImageKey, OTAImage, OTAImageHeader
 import zigpy.util
 
@@ -29,6 +29,7 @@ class Basic(zigpy.util.LocalLogMixin, ABC):
     REFRESH = datetime.timedelta(hours=12)
 
     def __init__(self):
+        self.config = {}
         self._cache = {}
         self._is_enabled = False
         self._locks = defaultdict(asyncio.Semaphore)
@@ -145,6 +146,7 @@ class Trådfri(Basic):
 
     async def initialize_provider(self, ota_config: Dict) -> None:
         self.info("OTA provider enabled")
+        self.config = ota_config
         await self.refresh_firmware_list()
         self.enable()
 
@@ -154,7 +156,8 @@ async def refresh_firmware_list(self) -> None:
 
         async with self._locks[LOCK_REFRESH]:
             async with aiohttp.ClientSession(headers=self.HEADERS) as req:
-                async with req.get(self.UPDATE_URL) as rsp:
+                url = self.config.get(CONF_OTA_IKEA_URL, self.UPDATE_URL)
+                async with req.get(url) as rsp:
                     # IKEA does not always respond with an appropriate Content-Type
                     # but the response is always JSON
                     if not (200 <= rsp.status <= 299):