1 files changed, 18 insertions, 11 deletions
diff --git a/test/unit/test_patterns_query_manager.py b/test/unit/test_patterns_query_manager.py
index 5daf3a0..4662e8a 100644
--- a/test/unit/test_patterns_query_manager.py
+++ b/test/unit/test_patterns_query_manager.py
@@ -20,6 +20,7 @@ Haketilo unit tests - building pattern tree and putting it in a content script
 import pytest
 import json
 from selenium.webdriver.support.ui import WebDriverWait
+from selenium.common.exceptions import TimeoutException
 
 from ..script_loader import load_script
 
@@ -240,25 +241,31 @@ def background_js():
 @pytest.mark.usefixtures('webextension')
 def test_pqm_script_injection(driver, execute_in_page):
     # Let's open a normal page in a second window. Window 0 will be used to make
-    # changed to IndexedDB and window 1 to test the working of content scripts.
+    # changes to IndexedDB and window 1 to test the working of content scripts.
     driver.execute_script('window.open("about:blank", "_blank");')
     WebDriverWait(driver, 10).until(lambda d: len(d.window_handles) == 2)
     windows = [*driver.window_handles]
 
-    def run_content_script():
-        driver.switch_to.window(windows[1])
-        driver.get('https://gotmyowndoma.in/index.html')
-        windows[1] = driver.window_handles[1]
+    def get_tree_json(driver):
         return driver.execute_script(
             '''
             return (document.getElementById("tree-json") || {}).innerText;
             ''')
 
-    for attempt in range(10):
+    def run_content_script():
+        driver.switch_to.window(windows[1])
+        driver.get('https://gotmyowndoma.in/index.html')
+        windows[1] = driver.current_window_handle
+        try:
+            return WebDriverWait(driver, 10).until(get_tree_json)
+        except TimeoutException:
+            pass
+
+    for attempt in range(2):
         json_txt = run_content_script()
-        if json.loads(json_txt) == {}:
+        if json_txt and json.loads(json_txt) == {}:
             break;
-        assert attempt != 9
+        assert attempt != 2
 
     driver.switch_to.window(windows[0])
     execute_in_page(load_script('common/indexeddb.js'))
@@ -271,12 +278,12 @@ def test_pqm_script_injection(driver, execute_in_page):
     }
     execute_in_page('returnval(save_items(arguments[0]));', sample_data)
 
-    for attempt in range(10):
-        tree_json = run_content_script()
+    for attempt in range(2):
+        tree_json = run_content_script() or '{}'
         json.loads(tree_json)
         if all([m['identifier'] in tree_json for m in sample_mappings]):
             break
-        assert attempt != 9
+        assert attempt != 2
 
     driver.switch_to.window(windows[0])
     execute_in_page(