From a0cf93a6be04309c157b4b8047538246ab41c196 Mon Sep 17 00:00:00 2001
From: Mirko Milovanovic <mir_ko@me.com>
Date: Fri, 21 Feb 2025 20:26:03 +0100
Subject: [PATCH] feat: enhance setupStats function with retry logic and
 improve error handling; update entrypoint to run main script as PID 1

---
 main.py                | 21 ++++++++++++---------
 monolith-entrypoint.sh |  6 ++++--
 2 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/main.py b/main.py
index 5a42bb3..f896cb3 100644
--- a/main.py
+++ b/main.py
@@ -181,19 +181,22 @@ def convert_to_bytes(down, downUnit):
     return float(down) * (1024 ** {'B': 0, 'KB': 1, 'MB': 2, 'GB': 3}[downUnit])
 
 @yaspin()
-def setupStats(driver: webdriver.Remote, url: str):
+def setupStats(driver: webdriver.Remote, url: str, retries: int = 5) -> webdriver.Remote:
     logger.log(logging.INFO, 'Setting up stats.')
     actions = ActionChains(driver)
     wait = WebDriverWait(driver, 30, poll_frequency=0.2)
 
-    driver.get(url)
-
-    try:
-        wait.until(ec.presence_of_element_located((By.CLASS_NAME, 'vjs-big-play-button')))
-    except Exception:
-        logger.error('Timeout while waiting for the big play button to be present.')
-        driver.quit()
-        raise SystemExit(1)
+    for attempt in range(retries):
+        driver.get(url)
+        try:
+            wait.until(ec.presence_of_element_located((By.CLASS_NAME, 'vjs-big-play-button')))
+            break
+        except Exception:
+            logger.error(f'Timeout while waiting for the big play button to be present. Attempt {attempt + 1} of {retries}')
+            if attempt == retries - 1:
+                logger.error('Timeout limit reached. Exiting.')
+                driver.quit()
+                raise SystemExit(1)
 
     actions.click(driver.find_element(By.CLASS_NAME ,'video-js')).perform()
     wait.until(ec.visibility_of_element_located((By.CLASS_NAME, 'vjs-control-bar')))
diff --git a/monolith-entrypoint.sh b/monolith-entrypoint.sh
index e02b4f0..65783b3 100644
--- a/monolith-entrypoint.sh
+++ b/monolith-entrypoint.sh
@@ -48,6 +48,8 @@ done
 
 printf '\n'
 
-# Start the Telegraf agent and the main script
+# Start the Telegraf agent
 telegraf --config ./telegraf.conf &
-./venv/bin/python main.py
\ No newline at end of file
+
+# Start the main Python script as PID 1
+exec ./venv/bin/python main.py
\ No newline at end of file