Skip to content

Commit

Permalink
fix: added test results and handled WebDriverException
Browse files Browse the repository at this point in the history
  • Loading branch information
@lynchee-owo
lynchee-owo committed May 15, 2024
1 parent 7f11954 commit 3513ee4
Show file tree
Hide file tree
Showing 208 changed files with 3,005 additions and 180 deletions.
153 changes: 96 additions & 57 deletions benchmark.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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):
Expand All @@ -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')
Expand All @@ -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()
Expand All @@ -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__":
Expand Down
Loading

0 comments on commit 3513ee4

Please sign in to comment.