testing actions

Author: jakubdotpy3d

.github/scripts/take_screenshot.py

@@ -1,14 +1,14 @@
+import time
+
 from selenium import webdriver
 from selenium.webdriver.chrome.options import Options
 from selenium.webdriver.common.by import By
 from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
 from PIL import Image
 import os
-import time
 
 
-# Initialize driver with options and cookie
 def setup_driver(cookie_value):
     options = Options()
     options.add_argument('--headless=new')
@@ -19,42 +19,31 @@ def setup_driver(cookie_value):
     return driver
 
 
-# Crop the screenshot using Pillow
 def crop_image(input_path, output_path, crop_box=(0, 0, 640, 621)):
     with Image.open(input_path) as img:
         cropped = img.crop(crop_box)
         cropped.save(output_path)
 
 
-# Take screenshot of an element and trim it
 def take_screenshot(driver, url, selector, output_name):
     driver.get(url)
-    time.sleep(3)  # Wait for the page to load completely
-    try:
-        # Wait for the element to be present and visible
-        element = WebDriverWait(driver, 20).until(
-            EC.presence_of_element_located((By.CSS_SELECTOR, selector))
-        )
-        temp_output = output_name  # Directly save the temp screenshot as the final output
-        element.screenshot(temp_output)
-        crop_image(temp_output, output_name)
-    except Exception as e:
-        print(f"Error: Timeout or issue locating element with selector '{selector}'")
-        print(f'Current URL: {driver.current_url}')
-        print('Page source:')
-        print(driver.page_source[:1000])  # Print first 1000 characters of source for debugging
-        raise
+    element = WebDriverWait(driver, 10).until(
+        EC.presence_of_element_located((By.CSS_SELECTOR, selector))
+    )
+    temp_output = output_name  # Directly save the temp screenshot as the final output
+    element.screenshot(temp_output)
+    crop_image(temp_output, output_name)
 
 
-# Main process
 def main():
     cookie = os.getenv('COOKIE')
-    if not cookie:
-        raise ValueError('COOKIE environment variable not set')
     driver = setup_driver(cookie)
     try:
         take_screenshot(
-            driver, 'https://adventofcode.com/2024', 'body > main > pre', 'aoc-screenshot.png'
+            driver,
+            'https://adventofcode.com/2024',
+            'body > main > pre',
+            'aoc-screenshot.png',
         )
     finally:
         driver.quit()