harness.py

Part of SNKRX.

  1#!/usr/bin/env python3
  2"""Scripted headless harness: drive the SNKRX port through a full run and
  3capture one screenshot per stage into ``screenshots/``.
  4
  5The driver starts the game, forces a slow-motion hit, then clears each arena
  6and confirms the shop so the run reaches the wave 5 boss inside 1400 frames.
  7
  8Stages
  9------
 101. Title screen  (frame 30)
 112. Wave 1 mid-combat  (frame 240)
 123. Slow-motion with screenshake  (frame 360)
 134. Shop after wave 1  (frame 540)
 145. Shop after wave 2  (frame 720)
 156. Shop after wave 3  (frame 900)
 167. Wave 4 in progress  (frame 1080)
 178. Wave 5 boss round  (frames 1260 and 1380)
 18
 19Run::
 20
 21    uv run python examples/ports/snkrx/harness.py
 22"""
 23
 24from __future__ import annotations
 25
 26import sys
 27from pathlib import Path
 28
 29_PORT_DIR = Path(__file__).resolve().parent
 30if str(_PORT_DIR) not in sys.path:
 31    sys.path.insert(0, str(_PORT_DIR))
 32
 33from main import WINDOW_H, WINDOW_W, SNKRXRoot  # noqa: E402
 34from nodes.colours import BG  # noqa: E402
 35
 36from simvx.core import InputSimulator, Key  # noqa: E402
 37from simvx.graphics import App, save_png  # noqa: E402
 38
 39
 40def _save(captures, frames, root_dir):
 41    root_dir.mkdir(exist_ok=True)
 42    for idx, frame in zip(frames, captures, strict=False):
 43        path = root_dir / f"stage_{idx:02d}.png"
 44        save_png(frame, path)
 45        print(f"saved {path}")
 46
 47
 48def _force_clear_arena(root):
 49    """Clear the active arena's enemies + queue so it ends next frame."""
 50    if root.phase == "arena" and root._sub is not None:
 51        arena = root._sub
 52        arena.spawn_queue.clear()
 53        for e in list(arena.enemies):
 54            e.alive = False
 55            e.hp = 0
 56
 57
 58def _drive(root: SNKRXRoot, sim: InputSimulator, frame: int):
 59    """Per-frame input driver: moves the harness through the staged scenes."""
 60    if frame == 60:
 61        sim.press_key(Key.ENTER)
 62    if frame == 62:
 63        sim.release_key(Key.ENTER)
 64    if frame == 360:
 65        # Force slow-mo
 66        if root.phase == "arena" and root._sub is not None:
 67            root._sub.start_slowmo()
 68    if frame == 480:
 69        _force_clear_arena(root)
 70    # Skip the buy screens automatically to advance to wave 5 boss
 71    if frame in (600, 800, 1000, 1200):
 72        if root.phase == "buy":
 73            sim.press_key(Key.ENTER)
 74    if frame in (602, 802, 1002, 1202):
 75        sim.release_key(Key.ENTER)
 76    if frame in (700, 900, 1100):
 77        _force_clear_arena(root)
 78
 79
 80def main():
 81    app = App(title="SNKRX Harness", width=WINDOW_W, height=WINDOW_H, visible=False, bg_colour=BG)
 82    root = SNKRXRoot()
 83    sim = InputSimulator()
 84
 85    captures_target = [30, 240, 360, 540, 720, 900, 1080, 1260, 1380]
 86
 87    def _per_frame(idx: int, _t: float):
 88        _drive(root, sim, idx)
 89        if idx in captures_target:
 90            print(f"frame {idx}: phase={root.phase}")
 91        return None  # continue
 92
 93    raw = app.run_headless(
 94        root,
 95        frames=1400,
 96        on_frame=_per_frame,
 97        capture_frames=captures_target,
 98    )
 99    _save(raw, captures_target, _PORT_DIR / "screenshots")
100
101
102if __name__ == "__main__":
103    main()