path: root/scripts/test_scripts/oplog-import-export-test.py

#!/usr/bin/env python3
"""Test script for oplog import/export operations."""

from __future__ import annotations

import argparse
import json
import os
import platform
import subprocess
import sys
from pathlib import Path
from typing import NamedTuple

_PLATFORM = "windows" if sys.platform == "win32" else "macosx" if sys.platform == "darwin" else "linux"
_ARCH = "x64" if sys.platform == "win32" else platform.machine().lower()
_EXE_SUFFIX = ".exe" if sys.platform == "win32" else ""


def _cache_dir() -> Path:
    if sys.platform == "win32":
        base = Path(os.environ.get("LOCALAPPDATA", Path.home() / "AppData" / "Local"))
        return base / "Temp" / "zen"
    elif sys.platform == "darwin":
        return Path.home() / "Library" / "Caches" / "zen"
    else:
        base = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache"))
        return base / "zen"


_BUILD_IDS_PATH = _cache_dir() / "oplog-import-export-build-ids.json"


class Build(NamedTuple):
    name: str
    bucket: str
    id: str


def load_builds() -> tuple[str, list[Build]]:
    if not _BUILD_IDS_PATH.exists():
        print(f"Build IDs file not found: {_BUILD_IDS_PATH}")
        answer = input("Run oplog-update-build-ids.py now to populate it? [y/N] ").strip().lower()
        if answer == "y":
            update_script = Path(__file__).parent / "oplog-update-build-ids.py"
            subprocess.run([sys.executable, str(update_script)], check=True)
        else:
            sys.exit("Aborted. Run scripts/test_scripts/oplog-update-build-ids.py to populate it.")
    with _BUILD_IDS_PATH.open() as f:
        data: dict = json.load(f)
    namespace = data.get("namespace", "")
    if not namespace:
        sys.exit(f"error: {_BUILD_IDS_PATH} is missing 'namespace'")
    builds = []
    for name, entry in data.get("builds", {}).items():
        bucket = entry.get("bucket", "")
        build_id = entry.get("buildId", "")
        if not bucket or not build_id:
            sys.exit(f"error: entry '{name}' in {_BUILD_IDS_PATH} is missing 'bucket' or 'buildId'")
        builds.append(Build(name, bucket, build_id))
    if not builds:
        sys.exit(f"error: {_BUILD_IDS_PATH} contains no builds")
    return namespace, builds


ZEN_EXE: Path = Path(f"./build/{_PLATFORM}/{_ARCH}/release/zen{_EXE_SUFFIX}")

ZEN_PORT = 8558
ZEN_CACHE_PORT = 8559
ZEN_CACHE = f"http://127.0.0.1:{ZEN_CACHE_PORT}"
ZEN_CACHE_POPULATE = "true"
ZEN_PARTIAL_REQUEST_MODE = "true"

SERVER_ARGS: tuple[str, ...] = (
    "--http", "asio",
    "--gc-cache-duration-seconds", "1209600",
    "--gc-interval-seconds", "21600",
    "--gc-low-diskspace-threshold", "2147483648",
    "--cache-bucket-limit-overwrites",
)


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:
    try:
        subprocess.run(cmd, check=True)
    except FileNotFoundError:
        sys.exit(f"error: executable not found: {cmd[0]}")
    except subprocess.CalledProcessError as e:
        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, 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}")
    try:
        subprocess.run(zen_cmd("down", "--port", str(port), extra_zen_args=extra_zen_args))
    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_zen_args: list[str] | None = None,
                 extra_server_args: list[str] | None = None) -> None:
    print(f"--------- starting {label} {data_dir}")
    run(zen_cmd(
        "up", "--port", str(port), "--show-console", "--",
        f"--data-dir={data_dir}",
        *SERVER_ARGS,
        *(extra_server_args or []),
        extra_zen_args=extra_zen_args,
    ))
    print()


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_cmd("wipe", "-y", path, extra_zen_args=extra_zen_args))
    else:
        print(f"--------- creating {label} {path}")
        path.mkdir(parents=True, exist_ok=True)
    print()


def check_prerequisites() -> None:
    if not ZEN_EXE.is_file():
        sys.exit(f"error: zen executable not found: {ZEN_EXE}")


def setup_project(port: int, extra_zen_args: list[str] | None = None) -> None:
    """Create the FortniteGame project on the server at the given port."""
    print("--------- creating FortniteGame project")
    run(zen_cmd("project-create", f"--hosturl=127.0.0.1:{port}", "FortniteGame", "--force-update",
                extra_zen_args=extra_zen_args))
    print()


def setup_oplog(port: int, build_name: str, extra_zen_args: list[str] | None = None) -> None:
    """Create the oplog in the FortniteGame project on the server at the given port."""
    print(f"--------- creating {build_name} oplog")
    run(zen_cmd("oplog-create", f"--hosturl=127.0.0.1:{port}", "FortniteGame", build_name, "--force-update",
                extra_zen_args=extra_zen_args))
    print()


def main() -> None:
    global ZEN_EXE

    # 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="?",
        default=None,
        type=Path,
        metavar="DATA_PATH",
        help="root path for all data directories (positional shorthand for --data-path)",
    )
    parser.add_argument(
        "zen_exe_positional",
        nargs="?",
        default=None,
        type=Path,
        metavar="ZEN_EXE_PATH",
        help="path to zen executable (positional shorthand for --zen-exe-path)",
    )
    parser.add_argument(
        "--data-path",
        default=None,
        type=Path,
        metavar="PATH",
        help="root path for all data directories",
    )
    parser.add_argument(
        "--zen-exe-path",
        default=ZEN_EXE,
        type=Path,
        metavar="PATH",
        help=f"path to zen executable (default: {ZEN_EXE})",
    )
    args = parser.parse_args(script_argv)

    data_path = args.positional_path
    if data_path is None:
        data_path = args.data_path
    if data_path is None:
        print("WARNING: This script may require up to 1TB of free disk space.")
        raw = input("Enter root path for all data directories: ").strip()
        if not raw:
            sys.exit("error: data path is required")
        data_path = Path(raw)

    ZEN_EXE = args.zen_exe_positional
    if ZEN_EXE is None:
        ZEN_EXE = args.zen_exe_path
    namespace, builds = load_builds()
    zen_data_dir = data_path / "DDC" / "OplogsZen"
    zen_cache_data_dir = data_path / "DDC" / "ZenBuildsCache"
    zen_import_data_dir = data_path / "DDC" / "OplogsZenImport"
    export_dir = data_path / "Export" / "FortniteGame"

    check_prerequisites()

    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("zenserver data", zen_data_dir, extra_zen_args)
        start_server("zenserver", zen_data_dir, ZEN_PORT, extra_zen_args=extra_zen_args)
        try:
            setup_project(ZEN_PORT, extra_zen_args)

            for build in builds:
                setup_oplog(ZEN_PORT, build.name, extra_zen_args)

                print(f"--------- importing {build.name} oplog")
                run(zen_cmd(
                    "oplog-import",
                    f"--hosturl=127.0.0.1:{ZEN_PORT}",
                    "FortniteGame", build.name,
                    "--clean",
                    "--builds", "https://jupiter.devtools.epicgames.com",
                    "--namespace", namespace,
                    "--bucket", build.bucket,
                    "--builds-id", build.id,
                    f"--zen-cache-host={ZEN_CACHE}",
                    f"--zen-cache-upload={ZEN_CACHE_POPULATE}",
                    f"--allow-partial-block-requests={ZEN_PARTIAL_REQUEST_MODE}",
                    extra_zen_args=extra_zen_args,
                ))
                print()

                print(f"--------- validating {build.name} oplog")
                run(zen_cmd("oplog-validate", f"--hosturl=127.0.0.1:{ZEN_PORT}", "FortniteGame", build.name,
                            extra_zen_args=extra_zen_args))
                print()

            wipe_or_create("export folder", export_dir, extra_zen_args)

            for build in builds:
                print(f"--------- exporting {build.name} oplog")
                run(zen_cmd(
                    "oplog-export",
                    f"--hosturl=127.0.0.1:{ZEN_PORT}",
                    "FortniteGame", build.name,
                    "--file", export_dir,
                    "--forcetempblocks",
                    extra_zen_args=extra_zen_args,
                ))
                print()
        finally:
            stop_server("zenserver", ZEN_PORT, extra_zen_args)

        wipe_or_create("alternate zenserver data", zen_import_data_dir, extra_zen_args)
        start_server("import zenserver", zen_import_data_dir, ZEN_PORT, extra_zen_args=extra_zen_args)
        try:
            setup_project(ZEN_PORT, extra_zen_args)

            for build in builds:
                setup_oplog(ZEN_PORT, build.name, extra_zen_args)

                print(f"--------- importing {build.name} oplog")
                run(zen_cmd(
                    "oplog-import",
                    f"--hosturl=127.0.0.1:{ZEN_PORT}",
                    "FortniteGame", build.name,
                    "--file", export_dir,
                    extra_zen_args=extra_zen_args,
                ))
                print()

                print(f"--------- validating {build.name} oplog")
                run(zen_cmd("oplog-validate", f"--hosturl=127.0.0.1:{ZEN_PORT}", "FortniteGame", build.name,
                            extra_zen_args=extra_zen_args))
                print()
        finally:
            stop_server("alternative zenserver", ZEN_PORT, extra_zen_args)
    finally:
        stop_server("cache zenserver", ZEN_CACHE_PORT, extra_zen_args)


if __name__ == "__main__":
    main()