ecm_sync_system/scripts/auto_test_init.py

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

from auto_test_lib import (
    create_couchdb_database,
    delete_couchdb_database,
    ensure_parent,
    load_yaml_config,
    recreate_dir,
    remove_path,
    run_command,
    start_background_process,
    stop_backend_processes,
    stop_process_from_pid_file,
    wait_http_ready,
)


def clear_environment(config_path: str | Path) -> dict:
    config = load_yaml_config(config_path)

    backend = config["backend"]
    couchdb = config.get("couchdb", {})
    paths = config["paths"]

    backend_stopped = stop_process_from_pid_file(backend["pid_file"])
    backend_killed_pgids = stop_backend_processes(backend["root_dir"])

    couchdb_delete_status = None
    if couchdb.get("enabled", False):
        couchdb_delete_status = delete_couchdb_database(
            couchdb["base_url"],
            couchdb["database"],
            couchdb.get("username", ""),
            couchdb.get("password", ""),
        )
        if couchdb.get("recreate_after_delete", False):
            create_couchdb_database(
                couchdb["base_url"],
                couchdb["database"],
                couchdb.get("username", ""),
                couchdb.get("password", ""),
            )

    persistence_db_existed = Path(paths["persistence_db"]).exists()
    temp_dir_existed = Path(paths["temp_data_dir"]).exists()
    pipeline_log_dir_existed = Path(paths["pipeline_log_dir"]).exists()

    remove_path(paths["persistence_db"])
    recreate_dir(paths["temp_data_dir"])
    recreate_dir(paths["pipeline_log_dir"])
    ensure_parent(backend["log_file"])

    return {
        "backend_stopped": backend_stopped,
        "backend_killed_pgids": backend_killed_pgids,
        "couchdb_delete_status": couchdb_delete_status,
        "persistence_db_removed": persistence_db_existed and not Path(paths["persistence_db"]).exists(),
        "temp_dir_reset": temp_dir_existed and Path(paths["temp_data_dir"]).exists(),
        "pipeline_log_dir_reset": pipeline_log_dir_existed and Path(paths["pipeline_log_dir"]).exists(),
        "persistence_db": paths["persistence_db"],
        "temp_data_dir": paths["temp_data_dir"],
    }


def initialize_environment(config_path: str | Path) -> dict:
    config = load_yaml_config(config_path)

    backend = config["backend"]
    couchdb = config.get("couchdb", {})
    paths = config["paths"]
    backend_env = backend.get("env", {})

    stop_process_from_pid_file(backend["pid_file"])
    stop_backend_processes(backend["root_dir"])

    if couchdb.get("enabled", False):
        delete_couchdb_database(
            couchdb["base_url"],
            couchdb["database"],
            couchdb.get("username", ""),
            couchdb.get("password", ""),
        )
        if couchdb.get("recreate_after_delete", False):
            create_couchdb_database(
                couchdb["base_url"],
                couchdb["database"],
                couchdb.get("username", ""),
                couchdb.get("password", ""),
            )

    remove_path(paths["persistence_db"])
    recreate_dir(paths["temp_data_dir"])
    recreate_dir(paths["pipeline_log_dir"])
    ensure_parent(backend["log_file"])
    ensure_parent(backend["pid_file"])

    restore_result = None
    for attempt in range(3):
        restore_result = run_command(backend["restore_command"], backend["root_dir"], extra_env=backend_env)
        if restore_result.returncode == 0:
            break
        if attempt < 2:
            time.sleep(2)
    if restore_result is None or restore_result.returncode != 0:
        stdout = "" if restore_result is None else restore_result.stdout
        stderr = "" if restore_result is None else restore_result.stderr
        raise RuntimeError(
            "restore_database 执行失败\n"
            f"stdout:\n{stdout}\n\n"
            f"stderr:\n{stderr}"
        )

    process = start_background_process(
        backend["start_command"],
        backend["root_dir"],
        backend["log_file"],
        extra_env=backend_env,
    )
    Path(backend["pid_file"]).write_text(str(process.pid), encoding="utf-8")

    wait_http_ready(
        backend["ready_url"],
        int(backend.get("ready_timeout_seconds", 120)),
        int(backend.get("ready_interval_seconds", 2)),
    )

    return {
        "restore_returncode": restore_result.returncode,
        "restore_stdout_tail": restore_result.stdout[-2000:],
        "restore_stderr_tail": restore_result.stderr[-2000:],
        "backend_pid": process.pid,
        "backend_log_file": backend["log_file"],
        "backend_root_dir": backend["root_dir"],
        "backend_ready_url": backend["ready_url"],
        "backend_start_command": backend["start_command"],
        "backend_restore_command": backend["restore_command"],
        "backend_env": backend_env,
        "couchdb": couchdb,
        "persistence_db": paths["persistence_db"],
        "temp_data_dir": paths["temp_data_dir"],
        "pipeline_log_dir": paths["pipeline_log_dir"],
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="初始化自动测试环境")
    parser.add_argument(
        "--config",
        default="tests/fixtures/auto_sync/auto_test_config.yaml",
        help="自动测试配置文件路径",
    )
    parser.add_argument(
        "--clear-only",
        action="store_true",
        help="只执行清空环境，不做恢复和启动",
    )
    args = parser.parse_args()

    if args.clear_only:
        result = clear_environment(args.config)
    else:
        result = initialize_environment(args.config)
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0


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