From a9b37ca24435de081b797e42fbf9afdb8a002dbb Mon Sep 17 00:00:00 2001
From: Paul Donald <newtwen+github@gmail.com>
Date: Sat, 4 May 2024 02:05:19 +0200
Subject: [PATCH] Test that app has ota and ota has listeners for OTA
 operations

Apparently in some situations they're absent. Let's guard against them.

Closes issue #225
---
 custom_components/zha_toolkit/ota.py | 40 ++++++++++++++++------------
 1 file changed, 23 insertions(+), 17 deletions(-)

diff --git a/custom_components/zha_toolkit/ota.py b/custom_components/zha_toolkit/ota.py
index 9fae296..073d2d4 100644
--- a/custom_components/zha_toolkit/ota.py
+++ b/custom_components/zha_toolkit/ota.py
@@ -158,28 +158,34 @@ async def download_sonoff_ota(listener, ota_dir):
 
 async def download_zigpy_ota(app, listener):
     LOGGER.debug("Zigpy download procedure starting")
-    for _, (ota, _) in app.ota._listeners.items():
-        if isinstance(ota, zigpy.ota.provider.FileStore):
-            # Skip files provider
-            continue
-        await ota.refresh_firmware_list()
-        for image_key, image in ota._cache.items():
-            url = getattr(image, "url", None)
-            LOGGER.error("Try getting %r, %r, %r", image_key, url, image)
-            try:
-                img = await app.ota.get_ota_image(
-                    image_key.manufacturer_id, image_key.image_type, model=None
-                )
-                LOGGER.info("Got image %r", getattr(img, "header", None))
-            except Exception as e:
-                LOGGER.error("%r while getting %r - %s", e, image_key, url)
+    if hasattr(app, 'ota') and hasattr(app.ota, '_listeners'):
+        for _, (ota, _) in app.ota._listeners.items():
+            if isinstance(ota, zigpy.ota.provider.FileStore):
+                # Skip files provider
+                continue
+            await ota.refresh_firmware_list()
+            for image_key, image in ota._cache.items():
+                url = getattr(image, "url", None)
+                LOGGER.error("Try getting %r, %r, %r", image_key, url, image)
+                try:
+                    img = await app.ota.get_ota_image(
+                        image_key.manufacturer_id, image_key.image_type, model=None
+                    )
+                    LOGGER.info("Got image %r", getattr(img, "header", None))
+                except Exception as e:
+                    LOGGER.error("%r while getting %r - %s", e, image_key, url)
+    else:
+        LOGGER.warning("Could not get ota object for download_zigpy_ota, try again")
 
 
 async def ota_update_images(
     app, listener, ieee, cmd, data, service, params, event_data
 ):
-    for _, (ota, _) in app.ota._listeners.items():
-        await ota.refresh_firmware_list()
+    if hasattr(app, 'ota') and hasattr(app.ota, '_listeners'):
+        for _, (ota, _) in app.ota._listeners.items():
+            await ota.refresh_firmware_list()
+    else:
+        LOGGER.warning("Could not get ota object for ota_update_images, try again")
 
 
 async def ota_notify(