Setup screenshots

selfdrive/ui/tests/test_ui/raylib_screenshots.py

@@ -4,11 +4,14 @@
 import shutil
 import time
 import pathlib
-from collections import namedtuple
+import subprocess
 
 import pyautogui
 import pywinctl
 
+from collections import namedtuple
+from collections.abc import Callable
+
 from cereal import log
 from cereal import messaging
 from cereal.messaging import PubMaster
@@ -219,37 +222,83 @@ def setup_onroad_full_alert_long_text(click, pm: PubMaster):
     time.sleep(0.05)
 
 
-CASES = {
+CASES: dict[str, Callable] = {
   "homescreen": setup_homescreen,
-  "homescreen_paired": setup_homescreen,
-  "homescreen_prime": setup_homescreen,
-  "homescreen_update_available": setup_homescreen_update_available,
-  "settings_device": setup_settings,
-  "settings_network": setup_settings_network,
-  "settings_network_advanced": setup_settings_network_advanced,
-  "settings_toggles": setup_settings_toggles,
-  "settings_software": setup_settings_software,
-  "settings_software_download": setup_settings_software_download,
-  "settings_software_release_notes": setup_settings_software_release_notes,
-  "settings_firehose": setup_settings_firehose,
-  "settings_developer": setup_settings_developer,
-  "keyboard": setup_keyboard,
-  "pair_device": setup_pair_device,
-  "offroad_alert": setup_offroad_alert,
-  "confirmation_dialog": setup_confirmation_dialog,
-  "experimental_mode_description": setup_experimental_mode_description,
-  "onroad": setup_onroad,
-  "onroad_sidebar": setup_onroad_sidebar,
-  "onroad_small_alert": setup_onroad_small_alert,
-  "onroad_medium_alert": setup_onroad_medium_alert,
-  "onroad_full_alert": setup_onroad_full_alert,
-  "onroad_full_alert_multiline": setup_onroad_full_alert_multiline,
-  "onroad_full_alert_long_text": setup_onroad_full_alert_long_text,
+  # "homescreen_paired": setup_homescreen,
+  # "homescreen_prime": setup_homescreen,
+  # "homescreen_update_available": setup_homescreen_update_available,
+  # "settings_device": setup_settings,
+  # "settings_network": setup_settings_network,
+  # "settings_network_advanced": setup_settings_network_advanced,
+  # "settings_toggles": setup_settings_toggles,
+  # "settings_software": setup_settings_software,
+  # "settings_software_download": setup_settings_software_download,
+  # "settings_software_release_notes": setup_settings_software_release_notes,
+  # "settings_firehose": setup_settings_firehose,
+  # "settings_developer": setup_settings_developer,
+  # "keyboard": setup_keyboard,
+  # "pair_device": setup_pair_device,
+  # "offroad_alert": setup_offroad_alert,
+  # "confirmation_dialog": setup_confirmation_dialog,
+  # "experimental_mode_description": setup_experimental_mode_description,
+  # "onroad": setup_onroad,
+  # "onroad_sidebar": setup_onroad_sidebar,
+  # "onroad_small_alert": setup_onroad_small_alert,
+  # "onroad_medium_alert": setup_onroad_medium_alert,
+  # "onroad_full_alert": setup_onroad_full_alert,
+  # "onroad_full_alert_multiline": setup_onroad_full_alert_multiline,
+  # "onroad_full_alert_long_text": setup_onroad_full_alert_long_text,
+}
+
+
+def fullscreen_click_primary_button(click, pm: PubMaster):
+  click(1600, 950)
+
+
+def fullscreen_click_secondary_button(click, pm: PubMaster):
+  click(500, 950)
+
+
+def software_setup_get_started_next(click, pm: PubMaster):
+  click(2000, 630)
+
+
+def software_setup_choose_software_click_openpilot(click, pm: PubMaster):
+  click(1200, 320)
+
+
+def software_setup_choose_software_click_custom(click, pm: PubMaster):
+  click(1200, 580)
+
+
+# These cases are for scripts that launch their own UI process (setup, updater, reset).
+# Each case is a list of additional steps to perform and screenshot (after initial screenshot).
+# Each item can also be a group of steps to do, with the screenshot at the end.
+SCRIPT_UI_CASES: dict[str, list | list[list]] = {
+  "setup": [
+    fullscreen_click_primary_button,  # Low voltage warning; click "Continue"
+    software_setup_get_started_next,  # Get started page; click arrow
+    [
+      # Do this in a group since we only want a screenshot of the warning
+      software_setup_choose_software_click_custom,  # Choose software page; click "Custom"
+      fullscreen_click_primary_button,  # Click "Continue"
+    ],
+    [fullscreen_click_secondary_button, software_setup_choose_software_click_openpilot],  # Go back to choose software page and click "openpilot"
+    [fullscreen_click_primary_button, lambda click, pm: time.sleep(1)],  # Click "Continue"; wait for networks to load
+    fullscreen_click_primary_button,  # "Download" button
+  ],
+  "updater": [],
+  "reset": [
+    fullscreen_click_primary_button,  # Click "Confirm" on initial confirmation
+    fullscreen_click_primary_button,  # Click "Confirm" on final warning
+  ],
 }
 
 
 class TestUI:
-  def __init__(self):
+  def __init__(self, window_title="UI"):
+    self.window_title = window_title
+
     os.environ["SCALE"] = os.getenv("SCALE", "1")
     sys.modules["mouseinfo"] = False
 
@@ -264,7 +313,7 @@ def setup(self):
       time.sleep(0.05)
     time.sleep(0.5)
     try:
-      self.ui = pywinctl.getWindowsWithTitle("UI")[0]
+      self.ui = pywinctl.getWindowsWithTitle(self.window_title)[0]
     except Exception as e:
       print(f"failed to find ui window, assuming that it's in the top left (for Xvfb) {e}")
       self.ui = namedtuple("bb", ["left", "top", "width", "height"])(0, 0, 2160, 1080)
@@ -280,13 +329,45 @@ def click(self, x: int, y: int, *args, **kwargs):
     pyautogui.mouseUp(self.ui.left + x, self.ui.top + y, *args, **kwargs)
 
   @with_processes(["ui"])
-  def test_ui(self, name, setup_case):
+  def test_ui(self, name, setup_case: Callable):
     self.setup()
     time.sleep(UI_DELAY)  # wait for UI to start
     setup_case(self.click, self.pm)
     self.screenshot(name)
 
 
+class TestScriptUI(TestUI):
+  def __init__(self, script_path: str, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    self._script_path = script_path
+    self._process = None
+
+  def __enter__(self):
+    self._process = subprocess.Popen([sys.executable, self._script_path])
+    return self
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    if self._process:
+      self._process.terminate()
+      try:
+        self._process.wait(timeout=5)
+      except subprocess.TimeoutExpired:
+        self._process.kill()
+      self._process = None
+
+  # Override the TestUI method to to run multiple tests, and to avoid starting another UI process
+  def test_ui(self, name, setup_cases: list[Callable] | list[list[Callable]]):
+    self.setup()
+    time.sleep(UI_DELAY)
+    self.screenshot(name)  # initial screenshot
+    # Run each setup case, taking a screenshot after each group
+    for i, case in enumerate(setup_cases):
+      group = case if isinstance(case, list) else [case]  # each case can be a single step or group of steps
+      for setup_case in group:
+        setup_case(self.click, self.pm)  # run each step in the group
+      self.screenshot(f"{name}_{i + 1}")  # take screenshot after each case group
+
+
 def create_screenshots():
   if TEST_OUTPUT_DIR.exists():
     shutil.rmtree(TEST_OUTPUT_DIR)
@@ -310,6 +391,12 @@ def create_screenshots():
 
       t.test_ui(name, setup)
 
+  for name, setup_cases in SCRIPT_UI_CASES.items():
+    with OpenpilotPrefix():
+      window_title = "System Reset" if name == "reset" else name.capitalize()
+      with TestScriptUI(f"system/ui/{name}.py", window_title) as launcher:
+        launcher.test_ui(name, setup_cases)
+
 
 if __name__ == "__main__":
   create_screenshots()