From c4e9a42c1a883b9e5b542fe03c2ef888893f7608 Mon Sep 17 00:00:00 2001
From: Sebastian Goscik <sebastian.goscik@live.co.uk>
Date: Sat, 8 Jul 2023 16:30:09 +0100
Subject: [PATCH] Block all calls to protect client when the connection is
 dropped and we are awaiting a reconnect

---
 unifi_protect_backup/downloader.py                | 3 +++
 unifi_protect_backup/event_listener.py            | 2 ++
 unifi_protect_backup/missing_event_checker.py     | 3 +++
 unifi_protect_backup/unifi_protect_backup_core.py | 5 +++++
 unifi_protect_backup/utils.py                     | 3 +++
 5 files changed, 16 insertions(+)

diff --git a/unifi_protect_backup/downloader.py b/unifi_protect_backup/downloader.py
index 67f6d22..67af7b9 100644
--- a/unifi_protect_backup/downloader.py
+++ b/unifi_protect_backup/downloader.py
@@ -82,6 +82,9 @@ class VideoDownloader:
         self.logger.info("Starting Downloader")
         while True:
             try:
+                # Wait for unifi protect to be connected
+                await self._protect.connect_event.wait()
+
                 event = await self.download_queue.get()
                 self.current_event = event
                 self.logger = logging.LoggerAdapter(self.base_logger, {'event': f' [{event.id}]'})
diff --git a/unifi_protect_backup/event_listener.py b/unifi_protect_backup/event_listener.py
index 5672ccc..66b928c 100644
--- a/unifi_protect_backup/event_listener.py
+++ b/unifi_protect_backup/event_listener.py
@@ -100,6 +100,7 @@ class EventListener:
         if self._protect.check_ws():
             logger.extra_debug("Websocket is connected.")
         else:
+            self._protect.connect_event.clear()
             logger.warning("Lost connection to Unifi Protect.")
 
             # Unsubscribe, close the session.
@@ -125,4 +126,5 @@ class EventListener:
                 # Back off for a little while
                 await asyncio.sleep(10)
 
+            self._protect.connect_event.set()
             logger.info("Re-established connection to Unifi Protect and to the websocket.")
diff --git a/unifi_protect_backup/missing_event_checker.py b/unifi_protect_backup/missing_event_checker.py
index b6fda33..cc850ce 100644
--- a/unifi_protect_backup/missing_event_checker.py
+++ b/unifi_protect_backup/missing_event_checker.py
@@ -123,6 +123,9 @@ class MissingEventChecker:
         logger.info("Starting Missing Event Checker")
         while True:
             try:
+                # Wait for unifi protect to be connected
+                await self._protect.connect_event.wait()
+
                 logger.extra_debug("Running check for missing events...")
 
                 wanted_events = await self._get_missing_events()
diff --git a/unifi_protect_backup/unifi_protect_backup_core.py b/unifi_protect_backup/unifi_protect_backup_core.py
index e41f44c..e153d66 100644
--- a/unifi_protect_backup/unifi_protect_backup_core.py
+++ b/unifi_protect_backup/unifi_protect_backup_core.py
@@ -193,6 +193,11 @@ class UnifiProtectBackup:
             else:
                 raise ConnectionError("Failed to connect to UniFi Protect after 10 attempts")
 
+            # Add a lock to the protect client that can be used to prevent code accessing the client when it has
+            # lost connection
+            self._protect.connect_event = asyncio.Event()
+            self._protect.connect_event.set()
+
             # Get a mapping of camera ids -> names
             logger.info("Found cameras:")
             for camera in self._protect.bootstrap.cameras.values():
diff --git a/unifi_protect_backup/utils.py b/unifi_protect_backup/utils.py
index 88a0cfd..d329953 100644
--- a/unifi_protect_backup/utils.py
+++ b/unifi_protect_backup/utils.py
@@ -286,6 +286,9 @@ async def get_camera_name(protect: ProtectApiClient, id: str):
 
     If the camera ID is not know, it tries refreshing the cached data
     """
+    # Wait for unifi protect to be connected
+    await protect.connect_event.wait()
+
     try:
         return protect.bootstrap.cameras[id].name
     except KeyError: