From e7978ee9ed418c7dbb79d9394c0c284fd4a5ef10 Mon Sep 17 00:00:00 2001
From: "Jonathan R. Madsen" <jonathanrmadsen@gmail.com>
Date: Wed, 10 Jan 2024 06:06:35 -0600
Subject: [PATCH] Update test_graphs.py

- new strategy for headless browser
---
 source/python/gui/source/tests/test_graphs.py | 35 ++++++++++++++++---
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/source/python/gui/source/tests/test_graphs.py b/source/python/gui/source/tests/test_graphs.py
index bc232fa9c..6b74037dc 100644
--- a/source/python/gui/source/tests/test_graphs.py
+++ b/source/python/gui/source/tests/test_graphs.py
@@ -1518,10 +1518,37 @@ def test_parse_uploaded_file():
 
 def set_up(ip_addr="localhost", ip_port="8051"):
     # works for linux, no browser pops up
-    fireFoxOptions = webdriver.FirefoxOptions()
-    fireFoxOptions.add_argument("--headless")
-    browser = webdriver.Firefox(options=fireFoxOptions)
-    browser.get("http://" + ip_addr + ":" + ip_port + "/")
+    webdriver_keys = webdriver.__dict__.keys()
+    browser = None
+    exception_vals = []
+
+    if "Chrome" in webdriver_keys and "ChromeOptions" in webdriver_keys:
+        try:
+            chromiumOptions = webdriver.ChromeOptions()
+            chromiumOptions.add_argument("--headless")
+            browser = webdriver.Chrome(options=chromiumOptions)
+        except Exception as e:
+            exception_vals.append(e)
+
+    if (
+        browser is None
+        and "Firefox" in webdriver_keys
+        and "FirefoxOptions" in webdriver_keys
+    ):
+        try:
+            fireFoxOptions = webdriver.FirefoxOptions()
+            fireFoxOptions.add_argument("--headless")
+            browser = webdriver.Firefox(options=fireFoxOptions)
+        except Exception as e:
+            exception_vals.append(e)
+
+    if browser:
+        browser.get("http://" + ip_addr + ":" + ip_port + "/")
+    elif exception_vals:
+        for itr in exception_vals:
+            print(f"Exception: {itr}")
+        raise exception_vals[0]
+
     return browser