iron/scripts/real_openlist_multi_backend_acceptance.py

from __future__ import annotations

import json
import os
import sys
from dataclasses import dataclass
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from scripts.real_openlist_acceptance import AcceptanceConfig, IronClient, read_openlist_webdav


def main() -> int:
    config = MultiBackendConfig.from_env()
    client = IronClient(config.base)

    print(f"[1/9] logging into Iron at {config.base.iron_base_url}")
    client.login()

    print("[2/9] ensuring multi-backend OpenList configuration")
    backends = ensure_backends(client, config)
    for backend in backends:
        check = client.check_backend(backend["id"])
        if check["status"] != "healthy":
            raise RuntimeError(f"backend {backend['name']} unhealthy: {json.dumps(check, ensure_ascii=False)}")
        print(f"       {backend['name']}: healthy at {check['detail']}")

    original_policy = next(policy for policy in client.list_policies() if policy["file_class"] == config.file_class)
    print(f"[3/9] saved original {config.file_class!r} policy")

    remote_backend_ids = [backend["id"] for backend in backends]
    try:
        print("[4/9] verifying two-remote replica policy")
        client.upsert_policy(
            config.file_class,
            require_local=True,
            stable_replica_count=2,
            opportunistic_replica_count=0,
            preferred_backend_ids=remote_backend_ids,
            excluded_backend_ids=[],
            max_non_local_size_bytes=None,
        )
        first_file_id = client.upload_text_file()
        first_detail = client.wait_for_ready_replica_set(first_file_id, {"bkd_local_default", remote_backend_ids[0], remote_backend_ids[1]})
        if config.verify_direct_webdav:
            assert_remote_bytes(config, backends, first_detail, remote_backend_ids[:2])
        print(f"       first file ready on local + {remote_backend_ids[0]} + {remote_backend_ids[1]}")

        print("[5/9] disabling the first remote backend")
        disabled_backend = client.disable_backend(remote_backend_ids[0])
        print(f"       disabled: {disabled_backend['name']}")

        print("[6/9] verifying new uploads avoid the disabled backend")
        second_filename = config.base.test_filename.replace(".txt", "-after-disable.txt")
        second_content = config.base.test_content.rstrip("\n") + "-after-disable\n"
        second_file_id = upload_named_text_file(client, config.base.directory_id, second_filename, second_content)
        second_detail = client.wait_for_ready_replica_set(second_file_id, {"bkd_local_default", remote_backend_ids[1], remote_backend_ids[2]})
        if config.verify_direct_webdav:
            assert_remote_bytes(config, backends, second_detail, remote_backend_ids[1:3], expected_text=second_content)
        print(f"       second file ready on local + {remote_backend_ids[1]} + {remote_backend_ids[2]}")

        print("[7/9] re-enabling the first remote backend")
        enabled_backend = client.enable_backend(remote_backend_ids[0])
        check = client.check_backend(enabled_backend["id"])
        if check["status"] != "healthy":
            raise RuntimeError(f"re-enabled backend not healthy: {json.dumps(check, ensure_ascii=False)}")
        print(f"       re-enabled: {enabled_backend['name']}")

        print("[8/9] verifying three-remote replica policy")
        client.upsert_policy(
            config.file_class,
            require_local=True,
            stable_replica_count=3,
            opportunistic_replica_count=0,
            preferred_backend_ids=remote_backend_ids,
            excluded_backend_ids=[],
            max_non_local_size_bytes=None,
        )
        third_filename = config.base.test_filename.replace(".txt", "-all-three.txt")
        third_content = config.base.test_content.rstrip("\n") + "-all-three\n"
        third_file_id = upload_named_text_file(client, config.base.directory_id, third_filename, third_content)
        third_detail = client.wait_for_ready_replica_set(third_file_id, {"bkd_local_default", *remote_backend_ids})
        if config.verify_direct_webdav:
            assert_remote_bytes(config, backends, third_detail, remote_backend_ids, expected_text=third_content)
        print("       third file ready on local + all remote backends")

    finally:
        print("[9/9] restoring original policy")
        client.upsert_policy(
            config.file_class,
            require_local=bool(original_policy["require_local"]),
            stable_replica_count=int(original_policy["stable_replica_count"]),
            opportunistic_replica_count=int(original_policy["opportunistic_replica_count"]),
            preferred_backend_ids=list(original_policy["preferred_backend_ids"]),
            excluded_backend_ids=list(original_policy["excluded_backend_ids"]),
            max_non_local_size_bytes=original_policy["max_non_local_size_bytes"],
        )
        print("       original policy restored")

    print("Multi-backend acceptance passed.")
    return 0


@dataclass
class MultiBackendConfig:
    base: AcceptanceConfig
    backend_names: list[str]
    root_paths: list[str]
    file_class: str
    verify_direct_webdav: bool

    @classmethod
    def from_env(cls) -> "MultiBackendConfig":
        base = AcceptanceConfig.from_env()
        raw_paths = os.getenv("OPENLIST_MULTI_ROOT_PATHS", "aliyun/iron,aliyun/iron-b,aliyun/iron-c")
        root_paths = [item.strip().strip("/") for item in raw_paths.split(",") if item.strip()]
        if len(root_paths) < 3:
            raise RuntimeError("OPENLIST_MULTI_ROOT_PATHS must define at least three root paths.")
        backend_names = [
            "openlist-aliyun",
            "openlist-aliyun-b",
            "openlist-aliyun-c",
            *[
                f"openlist-aliyun-{index}"
                for index in range(4, len(root_paths) + 1)
            ],
        ][: len(root_paths)]
        return cls(
            base=base,
            backend_names=backend_names,
            root_paths=root_paths,
            file_class=os.getenv("IRON_MULTI_FILE_CLASS", "other"),
            verify_direct_webdav=os.getenv("IRON_MULTI_VERIFY_DIRECT_WEBDAV", "").strip().lower() in {"1", "true", "yes", "on"},
        )


def ensure_backends(client: IronClient, config: MultiBackendConfig) -> list[dict[str, object]]:
    existing = {backend["name"]: backend for backend in client.list_backends()}
    ensured: list[dict[str, object]] = []
    for index, (name, root_path) in enumerate(zip(config.backend_names, config.root_paths, strict=True), start=1):
        backend = existing.get(name)
        if backend is None:
            payload = {
                "name": name,
                "type": "webdav",
                "stability_class": "stable",
                "read_priority": 90 - index,
                "write_priority": 90 - index,
                "config": {
                    "endpoint_url": config.base.openlist_dav_url,
                    "username": config.base.openlist_username,
                    "password": config.base.openlist_password,
                    "root_path": root_path,
                    "verify_ssl": False,
                },
            }
            backend = client._request("POST", "/api/backends", json_body=payload)["backend"]
        ensured.append(backend)
    return ensured


def upload_named_text_file(client: IronClient, directory_id: str, filename: str, content: str) -> str:
    original_filename = client.config.test_filename
    original_content = client.config.test_content
    try:
        client.config.test_filename = filename
        client.config.test_content = content
        return client.upload_text_file()
    finally:
        client.config.test_filename = original_filename
        client.config.test_content = original_content


def assert_remote_bytes(
    config: MultiBackendConfig,
    backends: list[dict[str, object]],
    detail: dict[str, object],
    backend_ids: list[str],
    *,
    expected_text: str | None = None,
) -> None:
    expected_bytes = (expected_text or config.base.test_content).encode("utf-8")
    replicas = detail["replicas"]
    backend_lookup = {backend["id"]: backend for backend in backends}
    for backend_id in backend_ids:
        replica = next(rep for rep in replicas if rep["backend_id"] == backend_id and rep["status"] == "ready")
        backend = backend_lookup[backend_id]
        try:
            body = read_openlist_webdav(
                endpoint_url=config.base.openlist_dav_url,
                username=config.base.openlist_username,
                password=config.base.openlist_password,
                relative_path=f"{backend['config']['root_path']}/{replica['storage_key']}",
            )
        except Exception as exc:
            raise RuntimeError(
                f"direct WebDAV read failed for backend {backend['name']} at "
                f"{backend['config']['root_path']}/{replica['storage_key']}: {exc}"
            ) from exc
        if body != expected_bytes:
            raise RuntimeError(f"unexpected bytes for backend {backend['name']}")


if __name__ == "__main__":
    raise SystemExit(main())