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())