fix: added test results and handled WebDriverException

Author: lynchee-owo

benchmark.py

@@ -8,14 +8,11 @@
 from selenium.webdriver.chrome.options import Options
 from selenium.webdriver.chrome.service import Service
 from webdriver_manager.chrome import ChromeDriverManager
+from selenium.common.exceptions import WebDriverException
 
 load_dotenv()
 api_key = os.getenv('OPENAI_API_KEY')
 
-# Hard-coded coordinates to open web-wand side panel
-extensions_pos = (1060, 110)
-web_wand_pos = (900, 280)
-
 # Place to store task execution results
 results_dir = 'results'
 os.makedirs(results_dir, exist_ok=True)
@@ -26,6 +23,8 @@ def setup_driver():
     chrome_options.add_argument("--load-extension=./dist")
     service = Service(ChromeDriverManager().install())
     driver = webdriver.Chrome(service=service, options=chrome_options)
+    # Set script timeout to 240 seconds
+    driver.set_script_timeout(240)
     return driver
 
 def dispatch_event(driver, event_name, event):
@@ -35,57 +34,84 @@ def dispatch_event(driver, event_name, event):
     """
     driver.execute_script(script)
 
-def add_task_listener(driver, task_id):
-    # Async script to add event listeners for taskStatus, taskHistory, and screenshot events
-    script = """
+def add_task_listener(driver, task_id, max_retries=3):
+    print('add_task_listener', task_id)
+    """
+    Add event listeners for task history and screenshot events. Both events include task status.
+    Then process those events as they are captured.
+    """
+
+    script = f"""
     var callback = arguments[0];
-    var keepListening = true;
-    var historyData = null;
-
-    var statusListener = function (e) {
-        if (e.detail.status !== 'running' && e.detail.status !== 'idle') {
-            keepListening = false;
-            document.removeEventListener('TaskStatusUpdate', statusListener);
-            document.removeEventListener('ScreenshotUpdate', screenshotListener);
-            document.removeEventListener('TaskHistoryUpdate', historyListener);
-            callback({type: 'status', data: 'stopped', history: historyData});
-        }
-    };
-
-    var screenshotListener = function (e) {
-        callback({type: 'screenshot', data: e.detail.imgData});
-    };
-
-    var historyListener = function (e) {
-        historyData = e.detail.history;
-    };
-
-    document.addEventListener('TaskStatusUpdate', statusListener);
-    document.addEventListener('ScreenshotUpdate', screenshotListener);
-    document.addEventListener('TaskHistoryUpdate', historyListener);
-
-    // To keep the async script alive
-    var checkInterval = setInterval(function () {
-        if (!keepListening) {
-            clearInterval(checkInterval);
-        }
-    }, 1000);
+    var eventListener = function (e) {{
+        if (e.detail.type == 'history') {{
+            console.log("event listener received history event");
+            if (e.detail.status === 'success' || e.detail.status === 'error') {{
+                callback({{status: e.detail.status, type: 'history', data: e.detail.data}});
+                document.removeEventListener('TaskUpdate', eventListener);
+                console.log("event listener removed");
+            }}
+            // Does not do anything when the status is 'running' or 'idle'. 
+            // The status 'interrupted' will never be triggered automatically.
+        }} else if (e.detail.type == 'screenshot') {{
+            console.log("event listener received screenshot event");
+            callback({{status: e.detail.status, type: 'screenshot', data: e.detail.data}});
+            document.removeEventListener('TaskUpdate', eventListener);
+            console.log("event listener removed");
+        }} else {{
+            throw new Error("Invalid event type received: " + e.detail.type);
+        }}
+    }};
+
+    document.addEventListener('TaskUpdate', eventListener);
+    console.log("added event listener");
     """
 
-    try:
-        while True:
-            event_data = driver.execute_async_script(script)
-            if event_data:
-                if event_data['type'] == 'screenshot':
-                    write_screenshots(task_id, event_data['data'])
-                elif event_data['type'] == 'status' and event_data['data'] == 'stopped':
-                    if event_data.get('history'):
-                        write_history(task_id, event_data['history'])
-                    break
-    except Exception as e:
-        print(f"Error while listening for updates: {e}")
+    completed = {'status': None}
+    attempts = 0
+
+    def handle_event(event_data):
+        nonlocal attempts
+        if not event_data:
+            print("no event data")
+            return
+        if event_data['type'] == 'history':
+            # Record history when task stops
+            completed['status'] = event_data['status']
+            write_history(task_id, event_data['data'])
+            return
+        if event_data['type'] == 'screenshot':
+            write_screenshots(task_id, event_data['data'])
+            # Task is still running. Continue to listen for events
+            handle_event(driver.execute_async_script(script))
+        else:
+            raise ValueError(f"Unhandled event data type: {event_data['type']}")
+        attempts = 0
+        print("reset attempts to zero")
+
+    while attempts < max_retries:
+        try:
+            handle_event(driver.execute_async_script(script))
+            break
+        except WebDriverException as e:
+            if "javascript error: document unloaded while waiting for result" in str(e):
+                print(f"Document unloaded error: {e}")
+                attempts += 1
+                print(f"Attempt {attempts}: Document unloaded error. Retrying...")
+                if attempts == max_retries:
+                    print("Maximum retry attempts reached. Cannot recover from document unloaded error.")
+            else:
+                print(f"Other WebDriver error: {e}")
+                break
+        except Exception as e:
+            print(f"Error while listening for updates: {e}")
+            break
+        
+    print("completed['status']", completed['status'])
+    return completed['status']
 
 def write_history(task_id, task_history):
+    print('write_history', task_id)
     task_dir = os.path.join(results_dir, f"test{task_id}")
     os.makedirs(task_dir, exist_ok=True)
     file_path = os.path.join(task_dir, 'interact_messages.json')
@@ -94,21 +120,33 @@ def write_history(task_id, task_history):
         json.dump(task_history, file, indent=4)
 
 def write_screenshots(task_id, image_data):
+    print('write_screenshots', task_id)
     image_bytes = base64.b64decode(image_data)
-
     task_dir = os.path.join(results_dir, f"test{task_id}")
     os.makedirs(task_dir, exist_ok=True)
     timestamp = int(time.time())
     file_path = os.path.join(task_dir, f'screenshot_{timestamp}.png')
-    
     with open(file_path, 'wb') as file:
         file.write(image_bytes)
 
 def run_webwand_task(driver, task_id, task_description):
+    print('run_webwand_task', task_id, task_description)
     dispatch_event(driver, 'SetAPIKey', {"value": api_key})
     dispatch_event(driver, 'SetTask', {"value": task_description})
     dispatch_event(driver, 'RunTask', {})
-    add_task_listener(driver, task_id)
+    task_status = add_task_listener(driver, task_id)
+    return task_status
+
+def click_extensions_icon(driver):
+    # Simulate click to open side panel
+    window_position = driver.get_window_rect()
+    top = window_position['y']
+    right = window_position['x'] + window_position['width']
+    # click Extensions icon
+    pyautogui.click(right - 150, top + 50)
+
+    # click webwand
+    pyautogui.click(right - 300, top + 210)
 
 def main():
     driver = setup_driver()
@@ -121,12 +159,13 @@ def main():
             driver.get(task['web'])
 
             if initial_load:
-                # Simulate click to open side panel
-                pyautogui.click(extensions_pos)
-                pyautogui.click(web_wand_pos)
+                click_extensions_icon(driver)
                 initial_load = False
 
-            run_webwand_task(driver, task_id, task['ques'])
+            task_status = run_webwand_task(driver, task_id, task['ques'])
+            while task_status not in ['success', 'error']:
+                print("wait task_status", task_status)
+                time.sleep(3)  # Wait for 3 seconds till the current task completes
     driver.quit()
 
 if __name__ == "__main__":