Fix AT-SPI crash and screenshot hangs, add screen capture after sleep

Framework/Built_In_Automation/Desktop/Linux/BuiltInFunctions.py

@@ -11,44 +11,116 @@
 
 from Framework.module_installer import install_missing_modules
 
+# ---------------------------------------------------------------------------
+# Optional system dependencies.
+#
+# pyatspi (AT-SPI2) is a thin Python layer over the GObject-Introspection
+# "Atspi" typelib, so the pip side (python3-pyatspi + pygobject) is only half of
+# it: the system typelib (gir1.2-atspi-2.0 / at-spi2-core) and a running
+# accessibility bus are required too. Having a display is not enough -- AT-SPI
+# is just as absent on a VNC/Xvfb desktop as on a bare server. python-xlib is
+# only used by the XComposite screenshot path, which already falls back to xwd.
+#
+# Importing this module therefore must never abort the process: CommonUtil
+# imports it just to take a desktop screenshot, and screenshot capture
+# (xdotool/xwd/XComposite based) does not need AT-SPI at all. A missing
+# dependency is recorded here and reported by the AT-SPI actions themselves as
+# an ordinary step failure.
+# ---------------------------------------------------------------------------
+
 try:
     import pyatspi
     from pyatspi.action import Action
     from pyatspi.editabletext import EditableText, Text
-except ImportError:
+
+    ATSPI_IMPORT_ERROR: Optional[str] = None
+except Exception:
     install_missing_modules(["python3-pyatspi==1.19.0", "pygobject==3.50.1"])
     try:
         import pyatspi
         from pyatspi.action import Action
         from pyatspi.editabletext import EditableText, Text
-    except ImportError:
-        sys.stderr.write(
-            "Error: system dependency is not installed. Install them by running Installer/setup_linux_inspector.sh.\n"
-        )
-        sys.exit(1)
+
+        ATSPI_IMPORT_ERROR = None
+    except Exception as atspi_error:
+        # pyatspi raises ImportError when the module is missing and ValueError
+        # ("Namespace Atspi not available") when the GI typelib is missing.
+        ATSPI_IMPORT_ERROR = f"{type(atspi_error).__name__}: {atspi_error}"
+
+        # The Accessible protocol below annotates methods with these interfaces,
+        # and annotations in a class body are evaluated at class-creation time,
+        # so the names must exist even when the real interfaces do not.
+        Action = EditableText = Text = Any  # type: ignore[misc,assignment]
+
+        class _AtspiUnavailable:
+            """Stand-in for the `pyatspi` module when AT-SPI is not installed.
+
+            Every attribute access raises, so an AT-SPI-dependent action fails as
+            a normal step failure with an actionable message (actions wrap their
+            body in try/except) instead of taking the node down.
+            """
+
+            def __getattr__(self, name: str):
+                raise RuntimeError(atspi_unavailable_message())
+
+        pyatspi = _AtspiUnavailable()  # type: ignore[assignment]
 
 try:
     from Xlib import X, display as xlib_display
     from Xlib.ext import composite as xlib_composite  # noqa: F401  registers ext methods
     from Xlib.error import XError
-except ImportError:
+
+    XLIB_IMPORT_ERROR: Optional[str] = None
+except Exception:
     install_missing_modules(["python-xlib==0.33"])
     try:
         from Xlib import X, display as xlib_display
         from Xlib.ext import composite as xlib_composite  # noqa: F401
         from Xlib.error import XError
-    except ImportError:
-        sys.stderr.write(
-            "Error: python-xlib is not installed. Install it by running Installer/setup_linux_inspector.sh.\n"
-        )
-        sys.exit(1)
+
+        XLIB_IMPORT_ERROR = None
+    except Exception as xlib_error:
+        XLIB_IMPORT_ERROR = f"{type(xlib_error).__name__}: {xlib_error}"
+        X = xlib_display = None  # type: ignore[assignment]
+
+        class XError(Exception):  # type: ignore[no-redef]
+            """Placeholder so `except XError` stays valid without python-xlib."""
 
 from Framework.Utilities import CommonUtil
 from Framework.Built_In_Automation.Shared_Resources import (
     BuiltInFunctionSharedResources as Shared_Resources,
 )
 from Framework.Utilities.decorators import logger
 
+_SETUP_HINT = "Installer/setup_linux_inspector.sh"
+
+
+def is_atspi_available() -> bool:
+    """True when the AT-SPI2 bindings imported successfully."""
+    return ATSPI_IMPORT_ERROR is None
+
+
+def atspi_unavailable_message() -> str:
+    """Actionable message for callers that need AT-SPI but cannot have it."""
+    return (
+        "Linux desktop automation requires AT-SPI2 (pyatspi), which is not importable on "
+        f"this machine ({ATSPI_IMPORT_ERROR}). Beyond the Python bindings it needs the system "
+        "AT-SPI2 typelib and a running accessibility bus -- on Debian/Ubuntu install "
+        "'at-spi2-core gir1.2-atspi-2.0' (or your distro's equivalent). "
+        f"See {_SETUP_HINT} for the X tools and accessibility settings."
+    )
+
+
+# Warn once, at import, rather than per action: the node keeps running, but
+# whoever reads the console should know why Linux desktop actions will fail.
+if ATSPI_IMPORT_ERROR:
+    sys.stderr.write(f"Warning: {atspi_unavailable_message()}\n")
+if XLIB_IMPORT_ERROR:
+    sys.stderr.write(
+        f"Warning: python-xlib is not available ({XLIB_IMPORT_ERROR}); "
+        f"screenshots fall back to xwd. See {_SETUP_HINT}.\n"
+    )
+
 
 class Collection: ...
 
@@ -328,6 +400,14 @@ def _capture_via_composite(file_path: str, winid: str) -> bool:
     Composite extension absent, deps missing, X errors) so the caller can
     fall back to a different strategy.
     """
+    if xlib_display is None:
+        CommonUtil.ExecLog(
+            MODULE_NAME,
+            f"Composite capture skipped: python-xlib unavailable ({XLIB_IMPORT_ERROR})",
+            4,
+        )
+        return False
+
     try:
         from PIL import Image
     except ImportError as e:
@@ -1013,7 +1093,13 @@ def _get_frame_geometry_for_window(window_id: str) -> dict | None:
     Used to align captured screenshots with the AT-SPI accessibility tree, whose
     coordinates may differ from the X11 window's geometry by the window
     decoration (title bar) offset.
+
+    Returns None when AT-SPI is unavailable, so the caller keeps the full
+    window pixmap instead of losing the screenshot entirely.
     """
+    if not is_atspi_available():
+        return None
+
     geometry = _get_window_geometry(window_id)
     if not geometry:
         return None

Framework/Built_In_Automation/Sequential_Actions/action_declarations/common.py

@@ -1,6 +1,6 @@
 declarations = (
     { "name": "step result",                                 "function": "step_result",                           "screenshot": "none" },
-    { "name": "sleep",                                       "function": "Sleep",                                 "screenshot": "none" },
+    { "name": "sleep",                                       "function": "Sleep",                                 "screenshot": "auto" },
     { "name": "wait",                                        "function": "Wait_For_Element",                      "screenshot": "none" },
     { "name": "wait disable",                                "function": "Wait_For_Element",                      "screenshot": "none" },
     { "name": "save text",                                   "function": "Save_Text",                             "screenshot": "none" },

Framework/Utilities/CommonUtil.py

@@ -3,6 +3,7 @@
 
 import selenium
 import sys
+import asyncio
 import inspect
 import os, os.path, threading
 import ast
@@ -821,6 +822,27 @@ def PhysicalAvailableMemory():
 )  # Initialize global variables for TakeScreenShot()
 
 
+AUTO_SCREEN_CAPTURE = "auto"
+
+
+def _resolve_auto_screen_capture(shared_variables):
+    """Pick the capture type for an action that is not tied to a platform.
+
+    Common actions (e.g. "sleep") are shared by every module, so their
+    declaration cannot name a platform up front. "auto" defers that choice to
+    here, where the drivers the test actually has open are visible.
+
+    Only web is resolved today: a sleep is most often used to let a page settle,
+    and the whole point of the capture is to show the page state once the wait
+    is over. Everything else stays "none" so no other module changes behaviour.
+    Add a "mobile" branch here if the same is ever wanted for Appium runs.
+    """
+    for driver_key in ("selenium_driver", "playwright_page"):
+        if shared_variables.get(driver_key) is not None:
+            return "web"
+    return "none"
+
+
 def set_screenshot_vars(shared_variables):
     """ Save screen capture type and selenium/appium driver objects as global variables, so TakeScreenShot() can access them """
     # We can't import Shared Variables due to cyclic imports causing local runs to break, so this is the work around
@@ -832,6 +854,8 @@ def set_screenshot_vars(shared_variables):
     try:
         if "screen_capture" in shared_variables:  # Type of screenshot (desktop/mobile)
             screen_capture_type = shared_variables["screen_capture"]
+        if screen_capture_type == AUTO_SCREEN_CAPTURE:
+            screen_capture_type = _resolve_auto_screen_capture(shared_variables)
         if screen_capture_type == "mobile":  # Appium driver object
             if "device_id" in shared_variables:
                 device_id = shared_variables[
@@ -1053,6 +1077,107 @@ def _screenshot_path(image_folder, image_name, extension="png"):
     return os.path.join(image_folder, safe_name + "." + extension.lstrip("."))
 
 
+_linux_capture_screenshot = None  # None = not resolved yet, False = resolution failed
+
+
+def _get_linux_capture_screenshot():
+    """Resolve the Linux desktop screenshot function once per process.
+
+    The Linux module pulls in optional system dependencies (AT-SPI, python-xlib)
+    that are absent on most nodes, and a module that fails to import is never
+    cached in sys.modules — so importing it per screenshot re-ran its dependency
+    install and re-raised per screenshot. Catch BaseException (it used to call
+    sys.exit(), and SystemExit slips straight past `except Exception`) and cache
+    the outcome, so a missing desktop stack costs one warning, not a traceback
+    on every capture.
+    """
+    global _linux_capture_screenshot
+
+    if _linux_capture_screenshot is None:
+        try:
+            from Framework.Built_In_Automation.Desktop.Linux.BuiltInFunctions import (
+                capture_screenshot as linux_capture_screenshot,
+            )
+
+            _linux_capture_screenshot = linux_capture_screenshot
+        except BaseException as e:
+            ExecLog(
+                MODULE_NAME,
+                "Linux desktop screenshot support is unavailable: %s: %s" % (type(e).__name__, e),
+                3,
+            )
+            _linux_capture_screenshot = False
+
+    return _linux_capture_screenshot or None
+
+
+SCREENSHOT_CAPTURE_TIMEOUT_SECONDS = 60
+SCREENSHOT_CAPTURE_POLL_SECONDS = 0.05
+
+
+def _log_capture_timeout(sModuleInfo, function_name, Method):
+    ExecLog(
+        sModuleInfo,
+        "Screenshot for Action: %s Method: %s did not finish within %s seconds and was "
+        "abandoned -- the browser/device is not responding. Continuing the run."
+        % (function_name, Method, SCREENSHOT_CAPTURE_TIMEOUT_SECONDS),
+        2,
+    )
+
+
+async def _capture_with_timeout(capture, sModuleInfo, function_name, Method):
+    """Run one screen capture bounded by a timeout, without blocking the loop.
+
+    Selenium and Appium screenshots are synchronous HTTP calls with no read
+    timeout of their own -- `RemoteConnection._timeout` defaults to
+    `socket._GLOBAL_DEFAULT_TIMEOUT`, so urllib3 waits forever -- and
+    `TakeScreenShot` runs outside `_run_action_with_timeout`. A wedged browser
+    therefore blocked the event loop indefinitely, parking the run on the
+    "Capturing Screenshot" log line with no timeout able to recover it.
+
+    The sync capture runs on a *daemon* thread and is simply abandoned on
+    timeout -- the same trade-off `_ActionTimeoutWorker` makes for a hung
+    action. It must be a daemon (and must not use the default executor, whose
+    threads `loop.shutdown_default_executor()` joins) so an abandoned capture
+    can never hold up node shutdown. Any exception is re-raised on this thread
+    so the existing handlers in `Thread_ScreenShot` still see it.
+
+    `capture` is a plain callable (Selenium/Appium/desktop) or an awaitable
+    (Playwright, already async). Returns True when it finished in time.
+    """
+    if inspect.isawaitable(capture):
+        try:
+            await asyncio.wait_for(capture, timeout=SCREENSHOT_CAPTURE_TIMEOUT_SECONDS)
+            return True
+        except (asyncio.TimeoutError, TimeoutError):
+            _log_capture_timeout(sModuleInfo, function_name, Method)
+            return False
+
+    done = threading.Event()
+    failure = {}
+
+    def runner():
+        try:
+            capture()
+        except BaseException:
+            failure["exc_info"] = sys.exc_info()
+        finally:
+            done.set()
+
+    threading.Thread(target=runner, name="zeuz_screenshot_capture", daemon=True).start()
+
+    deadline = time.monotonic() + SCREENSHOT_CAPTURE_TIMEOUT_SECONDS
+    while not done.is_set():
+        if time.monotonic() >= deadline:
+            _log_capture_timeout(sModuleInfo, function_name, Method)
+            return False
+        await asyncio.sleep(SCREENSHOT_CAPTURE_POLL_SECONDS)
+
+    if "exc_info" in failure:
+        raise failure["exc_info"][1].with_traceback(failure["exc_info"][2])
+    return True
+
+
 async def Thread_ScreenShot(function_name, image_folder, Method, Driver, image_name, skip_delay=False):
     """Capture screen of mobile, desktop, Selenium, or Playwright."""
     if performance_testing: return
@@ -1081,40 +1206,53 @@ async def Thread_ScreenShot(function_name, image_folder, Method, Driver, image_n
         if should_delay_before_capture and not skip_delay and not _wait_for_debug_screenshot_delay(sModuleInfo, function_name, Method):
             return
 
+        # Every capture below is bounded: a hung browser/device must not park the run.
+        def _desktop_grab():
+            image = ImageGrab_Mac_Win.grab(_get_window_screenshot_bbox())
+            image.save(ImageName, format="PNG")  # Save to disk
+
         # Capture screenshot of desktop
         if Method == "desktop":
             if sys.platform in ("linux", "linux2"):
-                # Import Linux screenshot function for AT-SPI desktop automation
-                try:
-                    if sys.platform in ("linux", "linux2"):
-                        from Framework.Built_In_Automation.Desktop.Linux.BuiltInFunctions import capture_screenshot as linux_capture_screenshot
-                except Exception:
-                    linux_capture_screenshot = None
-                if linux_capture_screenshot:
-                    linux_capture_screenshot(ImageName)
-                else:
+                linux_capture_screenshot = _get_linux_capture_screenshot()
+                if linux_capture_screenshot is None:
                     ExecLog(
                         sModuleInfo,
                         "Linux screenshot module not available",
                         3,
                     )
                     return
+                if not await _capture_with_timeout(
+                    lambda: linux_capture_screenshot(ImageName), sModuleInfo, function_name, Method
+                ):
+                    return
             elif sys.platform == "win32" or sys.platform == "darwin":
-                bbox = _get_window_screenshot_bbox()
-                image = ImageGrab_Mac_Win.grab(bbox)
-                image.save(ImageName, format="PNG")  # Save to disk
+                if not await _capture_with_timeout(_desktop_grab, sModuleInfo, function_name, Method):
+                    return
 
         # Capture screenshot of web browser
         elif Method == "web":
             # Check if it's a Playwright page or Selenium driver
             if is_playwright_page:
-                await Driver.screenshot(path=ImageName, type="jpeg", quality=PLAYWRIGHT_AUTO_SCREENSHOT_QUALITY)
+                captured = await _capture_with_timeout(
+                    Driver.screenshot(path=ImageName, type="jpeg", quality=PLAYWRIGHT_AUTO_SCREENSHOT_QUALITY),
+                    sModuleInfo, function_name, Method,
+                )
             else:  # Selenium driver
-                Driver.get_screenshot_as_file(ImageName)  # Must be .png, otherwise an exception occurs
+                # Must be .png, otherwise an exception occurs
+                captured = await _capture_with_timeout(
+                    lambda: Driver.get_screenshot_as_file(ImageName), sModuleInfo, function_name, Method
+                )
+            if not captured:
+                return
 
         # Capture screenshot of mobile
         elif Method == "mobile":
-            Driver.save_screenshot(ImageName)  # Must be .png, otherwise an exception occurs
+            # Must be .png, otherwise an exception occurs
+            if not await _capture_with_timeout(
+                lambda: Driver.save_screenshot(ImageName), sModuleInfo, function_name, Method
+            ):
+                return
         else:
             ExecLog(
                 sModuleInfo,