diff --git a/test/scripts/html_render_diff.py b/test/scripts/html_render_diff.py
index 5a601263..a945e78f 100755
--- a/test/scripts/html_render_diff.py
+++ b/test/scripts/html_render_diff.py
@@ -5,8 +5,13 @@
import sys
import argparse
import io
+import time
+
from PIL import Image, ImageChops
from selenium import webdriver
+from selenium.webdriver.common.by import By
+from selenium.webdriver.support import expected_conditions
+from selenium.webdriver.support.ui import WebDriverWait
import pathlib
@@ -19,8 +24,25 @@ def to_url(something):
def screenshot(browser, url):
browser.get(url)
- body = browser.find_element('tag name', 'body')
- png = body.screenshot_as_png
+ target_find_by = By.TAG_NAME
+ target = 'body'
+ loaded_page_settling_time = 0
+
+ # Selenium doesn't like when we try to screenshot element of documents generated by pdf2htmlEX
+ if 'output-pdf2htmlEX' in url:
+ target_find_by = By.ID
+ target = 'page-container'
+ loaded_page_settling_time = 1
+
+ web_driver_wait = WebDriverWait(browser, 5)
+ web_driver_wait.until(expected_conditions.presence_of_element_located((target_find_by, target)))
+ web_driver_wait.until(lambda driver: driver.execute_script("return document.readyState") == "complete")
+ if loaded_page_settling_time != 0:
+ time.sleep(loaded_page_settling_time)
+
+ target_element = browser.find_element(target_find_by, target)
+
+ png = target_element.screenshot_as_png
return Image.open(io.BytesIO(png))