From 0518d1ba48974686a2235e9ca9ec089e48b14b2b Mon Sep 17 00:00:00 2001
From: Mirko Milovanovic <mir_ko@me.com>
Date: Fri, 14 Feb 2025 15:15:59 +0100
Subject: [PATCH] feat: integrate yaspin for improved logging and error
 handling in Chrome driver setup

---
 main.py                               | 16 ++++++++++++----
 requirements.txt                      |  3 ++-
 webrtc-internals-exporter/override.js |  9 ++++-----
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/main.py b/main.py
index 479ef25..e0247ad 100644
--- a/main.py
+++ b/main.py
@@ -4,6 +4,7 @@ import time
 import socket
 import logging
 import os
+from yaspin import yaspin
 from functools import partial
 from http.server import HTTPServer
 from utils.PostHandler import Handler
@@ -34,6 +35,7 @@ def interrupt_handler(signum, driver: webdriver.Remote):
     driver.quit()
     raise SystemExit
 
+@yaspin()
 def setupChromeDriver():
     logger.log(logging.INFO, 'Setting up Chrome driver.')
     chrome_options = Options()
@@ -159,7 +161,8 @@ def downloadStats(driver: webdriver.Chrome, peersDict: dict):
 
 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):
     logger.log(logging.INFO, 'Setting up stats.')
     actions = ActionChains(driver)
@@ -167,7 +170,13 @@ def setupStats(driver: webdriver.Remote, url: str):
 
     driver.get(url)
 
-    wait.until(ec.presence_of_element_located((By.CLASS_NAME, 'vjs-big-play-button')))
+    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)
+
     actions.click(driver.find_element(By.CLASS_NAME ,'video-js')).perform()
     wait.until(ec.visibility_of_element_located((By.CLASS_NAME, 'vjs-control-bar')))
     actions.context_click(driver.find_element(By.CLASS_NAME ,'video-js')).perform()
@@ -193,7 +202,6 @@ if __name__ == '__main__':
     
     setupStats(driver, url)
     
-    logger.log(logging.INFO, 'Starting server collector.')
+    logger.info('Starting server collector.')
     httpd = HTTPServer(('', 9092), partial(Handler, downloadStats, driver, logger))
-    logger.info('Server collector started.')
     httpd.serve_forever()
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index dbaaec2..d23e83b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,3 @@
 selenium
-beautifulsoup4
\ No newline at end of file
+beautifulsoup4
+yaspin
\ No newline at end of file
diff --git a/webrtc-internals-exporter/override.js b/webrtc-internals-exporter/override.js
index 9139446..b945de5 100644
--- a/webrtc-internals-exporter/override.js
+++ b/webrtc-internals-exporter/override.js
@@ -22,7 +22,7 @@ class WebrtcInternalExporter {
     });
 
     window.postMessage({ event: "webrtc-internal-exporter:ready" });
-    this.collectAllStats();
+    setInterval(() => this.collectAndPostAllStats(), this.updateInterval);
   }
 
   randomId() {
@@ -41,7 +41,7 @@ class WebrtcInternalExporter {
     this.peerConnections.set(id, pc);
     pc.addEventListener("connectionstatechange", () => {
       log(`connectionStateChange: ${pc.connectionState}`);
-      this.collectAndPostSingleStat(id);
+      this.collectAndPostAllStats();
 
       if (pc.connectionState === "closed") {
         this.peerConnections.delete(id);
@@ -100,7 +100,7 @@ class WebrtcInternalExporter {
     log(`Single stat collected:`, [stats]);
   }
 
-  async collectAllStats() {
+  async collectAndPostAllStats() {
     const stats = [];
 
     for (const [id] of this.peerConnections) {
@@ -120,8 +120,7 @@ class WebrtcInternalExporter {
     );
 
     log(`Stats collected:`, stats);
-
-    setTimeout(this.collectAllStats.bind(this), this.updateInterval);
+    
     return stats;
   }