diff options
Diffstat (limited to 'scripts')
| -rw-r--r-- | scripts/test_scripts/builds-download-upload-test.py | 70 |
1 files changed, 48 insertions, 22 deletions
diff --git a/scripts/test_scripts/builds-download-upload-test.py b/scripts/test_scripts/builds-download-upload-test.py index f03528d98..e4fee7cb8 100644 --- a/scripts/test_scripts/builds-download-upload-test.py +++ b/scripts/test_scripts/builds-download-upload-test.py @@ -49,7 +49,13 @@ SERVER_ARGS: tuple[str, ...] = ( ) +def zen_cmd(*args: str | Path, extra_zen_args: list[str] | None = None) -> list[str | Path]: + """Build a zen CLI command list, inserting extra_zen_args before subcommands.""" + return [ZEN_EXE, *(extra_zen_args or []), *args] + + def run(cmd: list[str | Path]) -> None: + print(f" > {' '.join(str(x) for x in cmd)}") try: subprocess.run(cmd, check=True) except FileNotFoundError: @@ -58,31 +64,35 @@ def run(cmd: list[str | Path]) -> None: sys.exit(f"error: command failed with exit code {e.returncode}:\n {' '.join(str(x) for x in e.cmd)}") -def stop_server(label: str, port: int) -> None: +def stop_server(label: str, port: int, extra_zen_args: list[str] | None = None) -> None: """Stop a zen server. Tolerates failures so it is safe to call from finally blocks.""" print(f"--------- stopping {label}") + cmd = zen_cmd("down", "--port", str(port), extra_zen_args=extra_zen_args) + print(f" > {' '.join(str(x) for x in cmd)}") try: - subprocess.run([ZEN_EXE, "down", "--port", str(port)]) + subprocess.run(cmd) except OSError as e: print(f"warning: could not stop {label}: {e}", file=sys.stderr) print() -def start_server(label: str, data_dir: Path, port: int, extra_args: list[str] | None = None) -> None: +def start_server(label: str, data_dir: Path, port: int, extra_zen_args: list[str] | None = None, + extra_server_args: list[str] | None = None) -> None: print(f"--------- starting {label} {data_dir}") - run([ - ZEN_EXE, "up", "--port", str(port), "--show-console", "--", + run(zen_cmd( + "up", "--port", str(port), "--show-console", "--", f"--data-dir={data_dir}", *SERVER_ARGS, - *(extra_args or []), - ]) + *(extra_server_args or []), + extra_zen_args=extra_zen_args, + )) print() -def wipe_or_create(label: str, path: Path) -> None: +def wipe_or_create(label: str, path: Path, extra_zen_args: list[str] | None = None) -> None: if path.exists(): print(f"--------- cleaning {label} {path}") - run([ZEN_EXE, "wipe", "-y", path]) + run(zen_cmd("wipe", "-y", path, extra_zen_args=extra_zen_args)) else: print(f"--------- creating {label} {path}") path.mkdir(parents=True, exist_ok=True) @@ -103,7 +113,20 @@ def check_prerequisites() -> None: def main() -> None: global ZEN_EXE - parser = argparse.ArgumentParser(description=__doc__) + # Split on '--' to separate script args from extra zen CLI args + script_argv: list[str] = [] + extra_zen_args: list[str] = [] + if "--" in sys.argv[1:]: + sep = sys.argv.index("--", 1) + script_argv = sys.argv[1:sep] + extra_zen_args = sys.argv[sep + 1:] + else: + script_argv = sys.argv[1:] + + parser = argparse.ArgumentParser( + description=__doc__, + epilog="Any arguments after '--' are forwarded to every zen CLI invocation.", + ) parser.add_argument( "positional_path", nargs="?", @@ -134,7 +157,7 @@ def main() -> None: metavar="PATH", help=f"path to zen executable (default: {ZEN_EXE})", ) - args = parser.parse_args() + args = parser.parse_args(script_argv) data_path = args.positional_path if data_path is None: @@ -151,15 +174,16 @@ def main() -> None: check_prerequisites() - start_server("cache zenserver", zen_cache_data_dir, ZEN_CACHE_PORT, ["--buildstore-enabled"]) + start_server("cache zenserver", zen_cache_data_dir, ZEN_CACHE_PORT, + extra_zen_args=extra_zen_args, extra_server_args=["--buildstore-enabled"]) try: - wipe_or_create("download folder", zen_download_dir) - wipe_or_create("system folder", zen_system_dir) + wipe_or_create("download folder", zen_download_dir, extra_zen_args) + wipe_or_create("system folder", zen_system_dir, extra_zen_args) for build in BUILDS: print(f"--------- importing {build.name} build") - run([ - ZEN_EXE, "builds", "download", + run(zen_cmd( + "builds", "download", "--host", "https://jupiter.devtools.epicgames.com", "--namespace", "fortnite.oplog", "--bucket", build.bucket, @@ -169,15 +193,16 @@ def main() -> None: f"--allow-partial-block-requests={ZEN_PARTIAL_REQUEST_MODE}", "--verify", "--system-dir", zen_system_dir, - ]) + extra_zen_args=extra_zen_args, + )) print() - wipe_or_create("upload folder", zen_upload_dir) + wipe_or_create("upload folder", zen_upload_dir, extra_zen_args) for build in BUILDS: print(f"--------- exporting {build.name} build") - run([ - ZEN_EXE, "builds", "upload", + run(zen_cmd( + "builds", "upload", "--storage-path", zen_upload_dir, "--build-id", build.id, "--local-path", zen_download_dir / build.name, @@ -186,10 +211,11 @@ def main() -> None: "--metadata-path", str(ZEN_METADATA_DIR / f"{build.name}.json"), "--create-build", "--chunking-cache-path", zen_chunk_cache_path, - ]) + extra_zen_args=extra_zen_args, + )) print() finally: - stop_server("cache zenserver", ZEN_CACHE_PORT) + stop_server("cache zenserver", ZEN_CACHE_PORT, extra_zen_args) if __name__ == "__main__": |