mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
feat(cli): enhance backup capabilities and refactor path management
- Introduce 'astrbot bk' command with GPG signing, encryption, and digest support - Add import/export functionality using core backup modules - Refactor path management to use 'AstrbotPaths' singleton across CLI commands - Replace blocking subprocess calls with asyncio.create_subprocess_exec in backup command - Add comprehensive tests for uninstall and backup commands - Improve module resource handling for bundled dashboard assets
This commit is contained in:
294
astrbot/cli/commands/cmd_bk.py
Normal file
294
astrbot/cli/commands/cmd_bk.py
Normal file
@@ -0,0 +1,294 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import shutil
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
import click
|
||||
|
||||
from astrbot.core import astrbot_config, db_helper
|
||||
from astrbot.core.backup import AstrBotExporter, AstrBotImporter
|
||||
|
||||
# Try importing KnowledgeBaseManager to support KB backup
|
||||
try:
|
||||
from astrbot.core.knowledge.kb_manager import KnowledgeBaseManager
|
||||
except ImportError:
|
||||
try:
|
||||
from astrbot.core.knowledge_base.kb_manager import KnowledgeBaseManager
|
||||
except ImportError:
|
||||
KnowledgeBaseManager = None
|
||||
|
||||
|
||||
async def _get_kb_manager():
|
||||
if KnowledgeBaseManager is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
# Best effort initialization
|
||||
kb_mgr = KnowledgeBaseManager(astrbot_config, db_helper)
|
||||
# If there are async load methods, we might need to call them
|
||||
if hasattr(kb_mgr, "load_kbs_from_db"):
|
||||
await kb_mgr.load_kbs_from_db()
|
||||
elif hasattr(kb_mgr, "load_all"):
|
||||
await kb_mgr.load_all()
|
||||
return kb_mgr
|
||||
except Exception:
|
||||
# If KB manager fails to load (e.g. missing dependencies), return None
|
||||
# so we can still backup other data
|
||||
return None
|
||||
|
||||
|
||||
@click.group(name="bk")
|
||||
def bk():
|
||||
"""Backup management (Export/Import)"""
|
||||
pass
|
||||
|
||||
|
||||
@bk.command(name="export")
|
||||
@click.option("--output", "-o", help="Output directory", default=None)
|
||||
@click.option(
|
||||
"--gpg-sign", "-S", is_flag=True, help="Sign backup with GPG default private key"
|
||||
)
|
||||
@click.option(
|
||||
"--gpg-encrypt",
|
||||
"-E",
|
||||
help="Encrypt for GPG recipient (Asymmetric)",
|
||||
metavar="RECIPIENT",
|
||||
)
|
||||
@click.option(
|
||||
"--gpg-symmetric", "-C", is_flag=True, help="Encrypt with symmetric cipher (GPG)"
|
||||
)
|
||||
@click.option(
|
||||
"--digest",
|
||||
"-d",
|
||||
type=click.Choice(["md5", "sha1", "sha256", "sha512"]),
|
||||
help="Generate digital digest",
|
||||
)
|
||||
def export_data(
|
||||
output: str | None,
|
||||
gpg_sign: bool,
|
||||
gpg_encrypt: str | None,
|
||||
gpg_symmetric: bool,
|
||||
digest: str | None,
|
||||
):
|
||||
"""Export all AstrBot data to a backup archive.
|
||||
|
||||
If any GPG option (-S, -E, -C) is used, the output file will be processed by GPG
|
||||
and saved with a .gpg extension.
|
||||
|
||||
Examples:
|
||||
|
||||
\b
|
||||
1. Standard Export:
|
||||
astrbot bk export
|
||||
-> Generates a plain .zip file.
|
||||
|
||||
\b
|
||||
2. Signed Backup (Integrity Check):
|
||||
astrbot bk export -S
|
||||
-> Generates a .zip.gpg file containing the backup and your signature.
|
||||
-> NOT ENCRYPTED, but packaged in OpenPGP format.
|
||||
-> Use 'astrbot bk import' or 'gpg --verify' to check integrity.
|
||||
|
||||
\b
|
||||
3. Password Protected (Symmetric Encryption):
|
||||
astrbot bk export -C
|
||||
-> Generates an encrypted .zip.gpg file.
|
||||
-> Prompts for a passphrase.
|
||||
-> Only accessible with the passphrase.
|
||||
|
||||
\b
|
||||
4. Encrypted for Recipient (Asymmetric Encryption):
|
||||
astrbot bk export -E "alice@example.com"
|
||||
-> Generates an encrypted .zip.gpg file for Alice.
|
||||
-> Only Alice's private key can decrypt it.
|
||||
|
||||
\b
|
||||
5. Signed and Encrypted with Digest:
|
||||
astrbot bk export -S -E "bob@example.com" -d sha256
|
||||
-> Signs, encrypts for Bob, and generates a SHA256 checksum file.
|
||||
"""
|
||||
|
||||
async def _run():
|
||||
if gpg_sign or gpg_encrypt or gpg_symmetric:
|
||||
if not shutil.which("gpg"):
|
||||
raise click.ClickException(
|
||||
"GPG tool not found. Please install GnuPG to use encryption/signing features."
|
||||
)
|
||||
|
||||
kb_mgr = await _get_kb_manager()
|
||||
exporter = AstrBotExporter(db_helper, kb_mgr)
|
||||
|
||||
async def on_progress(stage, current, total, message):
|
||||
click.echo(f"[{stage}] {message}")
|
||||
|
||||
try:
|
||||
path_str = await exporter.export_all(output, progress_callback=on_progress)
|
||||
final_path = Path(path_str)
|
||||
click.echo(
|
||||
click.style(f"\nRaw backup exported to: {final_path}", fg="green")
|
||||
)
|
||||
|
||||
# GPG Operations
|
||||
if gpg_sign or gpg_encrypt or gpg_symmetric:
|
||||
# Construct GPG command
|
||||
# output file usually ends with .gpg
|
||||
gpg_output = final_path.with_name(final_path.name + ".gpg")
|
||||
cmd = ["gpg", "--output", str(gpg_output), "--yes"]
|
||||
|
||||
if gpg_symmetric:
|
||||
if gpg_encrypt:
|
||||
click.echo(
|
||||
click.style(
|
||||
"Warning: Symmetric encryption selected, ignoring asymmetric recipient.",
|
||||
fg="yellow",
|
||||
)
|
||||
)
|
||||
cmd.append("--symmetric")
|
||||
# No --batch to allow interactive passphrase entry on TTY
|
||||
else:
|
||||
# Asymmetric or just Sign
|
||||
# Note: If encrypting, -s adds signature to the encrypted packet.
|
||||
if gpg_encrypt:
|
||||
cmd.extend(["--encrypt", "--recipient", gpg_encrypt])
|
||||
|
||||
if gpg_sign:
|
||||
cmd.append("--sign")
|
||||
|
||||
cmd.append(str(final_path))
|
||||
|
||||
click.echo(f"Running GPG: {' '.join(cmd)}")
|
||||
|
||||
# Replace subprocess.run with asyncio.create_subprocess_exec to avoid blocking the event loop
|
||||
process = await asyncio.create_subprocess_exec(*cmd)
|
||||
await process.wait()
|
||||
|
||||
if process.returncode != 0:
|
||||
raise subprocess.CalledProcessError(process.returncode or 1, cmd)
|
||||
|
||||
# Clean up original file
|
||||
final_path.unlink()
|
||||
final_path = gpg_output
|
||||
click.echo(
|
||||
click.style(f"Processed backup created: {final_path}", fg="green")
|
||||
)
|
||||
|
||||
# Digest Generation
|
||||
if digest:
|
||||
click.echo(f"Calculating {digest} digest...")
|
||||
hash_func = getattr(hashlib, digest)()
|
||||
# Read file in chunks
|
||||
with open(final_path, "rb") as f:
|
||||
while chunk := f.read(8192):
|
||||
hash_func.update(chunk)
|
||||
|
||||
digest_val = hash_func.hexdigest()
|
||||
digest_file = final_path.with_name(final_path.name + f".{digest}")
|
||||
digest_file.write_text(
|
||||
f"{digest_val} *{final_path.name}\n", encoding="utf-8"
|
||||
)
|
||||
click.echo(click.style(f"Digest generated: {digest_file}", fg="green"))
|
||||
|
||||
except subprocess.CalledProcessError as e:
|
||||
click.echo(click.style(f"\nGPG process failed: {e}", fg="red"), err=True)
|
||||
except Exception as e:
|
||||
click.echo(click.style(f"\nExport failed: {e}", fg="red"), err=True)
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
@bk.command(name="import")
|
||||
@click.argument("backup_file")
|
||||
@click.option("--yes", "-y", is_flag=True, help="Skip confirmation prompts")
|
||||
def import_data(backup_file: str, yes: bool):
|
||||
"""Import AstrBot data from a backup archive.
|
||||
|
||||
Automatically handles .zip files and .gpg files (signed or encrypted).
|
||||
If the file is encrypted, you will be prompted for the passphrase.
|
||||
"""
|
||||
backup_path = Path(backup_file)
|
||||
if not backup_path.exists():
|
||||
raise click.ClickException(f"Backup file not found: {backup_file}")
|
||||
|
||||
if not yes:
|
||||
click.confirm(
|
||||
"This will OVERWRITE all current data (DB, Config, Plugins). Continue?",
|
||||
abort=True,
|
||||
default=False,
|
||||
)
|
||||
|
||||
async def _run():
|
||||
zip_path = backup_path
|
||||
is_temp_file = False
|
||||
|
||||
# Handle GPG encrypted files
|
||||
if backup_path.suffix == ".gpg":
|
||||
if not shutil.which("gpg"):
|
||||
raise click.ClickException(
|
||||
"GPG tool not found. Cannot decrypt .gpg file."
|
||||
)
|
||||
|
||||
# Remove .gpg extension for output
|
||||
decrypted_path = backup_path.with_suffix("")
|
||||
# If it doesn't look like a zip after stripping .gpg, maybe append .zip?
|
||||
# But the exporter creates .zip.gpg, so stripping .gpg gives .zip.
|
||||
|
||||
click.echo(f"Processing GPG file {backup_path}...")
|
||||
try:
|
||||
cmd = [
|
||||
"gpg",
|
||||
"--output",
|
||||
str(decrypted_path),
|
||||
"--decrypt", # This handles both decryption and signature verification/extraction
|
||||
str(backup_path),
|
||||
]
|
||||
# Allow interactive passphrase
|
||||
process = await asyncio.create_subprocess_exec(*cmd)
|
||||
await process.wait()
|
||||
|
||||
if process.returncode != 0:
|
||||
raise subprocess.CalledProcessError(process.returncode or 1, cmd)
|
||||
|
||||
zip_path = decrypted_path
|
||||
is_temp_file = True
|
||||
except subprocess.CalledProcessError:
|
||||
click.echo(
|
||||
click.style(
|
||||
"GPG processing failed. Verify signature or decryption key.",
|
||||
fg="red",
|
||||
),
|
||||
err=True,
|
||||
)
|
||||
return
|
||||
|
||||
kb_mgr = await _get_kb_manager()
|
||||
importer = AstrBotImporter(db_helper, kb_mgr)
|
||||
|
||||
async def on_progress(stage, current, total, message):
|
||||
click.echo(f"[{stage}] {message}")
|
||||
|
||||
try:
|
||||
result = await importer.import_all(
|
||||
str(zip_path), progress_callback=on_progress
|
||||
)
|
||||
|
||||
if result.errors:
|
||||
click.echo(
|
||||
click.style("\nImport failed with errors:", fg="red"), err=True
|
||||
)
|
||||
for err in result.errors:
|
||||
click.echo(f" - {err}", err=True)
|
||||
else:
|
||||
click.echo(click.style("\nImport completed successfully!", fg="green"))
|
||||
|
||||
if result.warnings:
|
||||
click.echo(click.style("\nWarnings:", fg="yellow"))
|
||||
for warn in result.warnings:
|
||||
click.echo(f" - {warn}")
|
||||
|
||||
finally:
|
||||
if is_temp_file and zip_path.exists():
|
||||
zip_path.unlink()
|
||||
click.echo(f"Cleaned up temporary file: {zip_path}")
|
||||
|
||||
asyncio.run(_run())
|
||||
@@ -6,7 +6,9 @@ from typing import Any
|
||||
|
||||
import click
|
||||
|
||||
from ..utils import check_astrbot_root, get_astrbot_root
|
||||
from astrbot.core.utils.astrbot_path import astrbot_paths
|
||||
|
||||
from ..utils import check_astrbot_root
|
||||
|
||||
|
||||
def _validate_log_level(value: str) -> str:
|
||||
@@ -77,13 +79,13 @@ CONFIG_VALIDATORS: dict[str, Callable[[str], Any]] = {
|
||||
|
||||
def _load_config() -> dict[str, Any]:
|
||||
"""Load or initialize config file"""
|
||||
root = get_astrbot_root()
|
||||
root = astrbot_paths.root
|
||||
if not check_astrbot_root(root):
|
||||
raise click.ClickException(
|
||||
f"{root} is not a valid AstrBot root directory. Use 'astrbot init' to initialize",
|
||||
)
|
||||
|
||||
config_path = root / "data" / "cmd_config.json"
|
||||
config_path = astrbot_paths.data / "cmd_config.json"
|
||||
if not config_path.exists():
|
||||
from astrbot.core.config.default import DEFAULT_CONFIG
|
||||
|
||||
@@ -100,7 +102,7 @@ def _load_config() -> dict[str, Any]:
|
||||
|
||||
def _save_config(config: dict[str, Any]) -> None:
|
||||
"""Save config file"""
|
||||
config_path = get_astrbot_root() / "data" / "cmd_config.json"
|
||||
config_path = astrbot_paths.data / "cmd_config.json"
|
||||
|
||||
config_path.write_text(
|
||||
json.dumps(config, ensure_ascii=False, indent=2),
|
||||
|
||||
@@ -7,7 +7,9 @@ from pathlib import Path
|
||||
import click
|
||||
from filelock import FileLock, Timeout
|
||||
|
||||
from ..utils import check_dashboard, get_astrbot_root
|
||||
from astrbot.core.utils.astrbot_path import astrbot_paths
|
||||
|
||||
from ..utils import check_dashboard
|
||||
|
||||
SYSTEMD_SERVICE = r"""
|
||||
# user service
|
||||
@@ -97,7 +99,7 @@ def init(yes: bool) -> None:
|
||||
except subprocess.CalledProcessError as e:
|
||||
click.echo(f"Failed to reload systemd daemon: {e}", err=True)
|
||||
|
||||
astrbot_root = get_astrbot_root()
|
||||
astrbot_root = astrbot_paths.root
|
||||
lock_file = astrbot_root / "astrbot.lock"
|
||||
lock = FileLock(lock_file, timeout=5)
|
||||
|
||||
|
||||
@@ -4,11 +4,12 @@ from pathlib import Path
|
||||
|
||||
import click
|
||||
|
||||
from astrbot.core.utils.astrbot_path import astrbot_paths
|
||||
|
||||
from ..utils import (
|
||||
PluginStatus,
|
||||
build_plug_list,
|
||||
check_astrbot_root,
|
||||
get_astrbot_root,
|
||||
get_git_repo,
|
||||
manage_plugin,
|
||||
)
|
||||
@@ -20,12 +21,12 @@ def plug() -> None:
|
||||
|
||||
|
||||
def _get_data_path() -> Path:
|
||||
base = get_astrbot_root()
|
||||
base = astrbot_paths.root
|
||||
if not check_astrbot_root(base):
|
||||
raise click.ClickException(
|
||||
f"{base} is not a valid AstrBot root directory. Use 'astrbot init' to initialize",
|
||||
)
|
||||
return (base / "data").resolve()
|
||||
return astrbot_paths.data.resolve()
|
||||
|
||||
|
||||
def display_plugins(plugins, title=None, color=None) -> None:
|
||||
|
||||
@@ -7,7 +7,9 @@ from pathlib import Path
|
||||
import click
|
||||
from filelock import FileLock, Timeout
|
||||
|
||||
from ..utils import check_astrbot_root, check_dashboard, get_astrbot_root
|
||||
from astrbot.core.utils.astrbot_path import astrbot_paths
|
||||
|
||||
from ..utils import check_astrbot_root, check_dashboard
|
||||
|
||||
|
||||
async def run_astrbot(astrbot_root: Path) -> None:
|
||||
@@ -51,7 +53,7 @@ def run(reload: bool, host: str, port: str, backend_only: bool, log_level: str)
|
||||
"""Run AstrBot"""
|
||||
try:
|
||||
os.environ["ASTRBOT_CLI"] = "1"
|
||||
astrbot_root = get_astrbot_root()
|
||||
astrbot_root = astrbot_paths.root
|
||||
|
||||
if not check_astrbot_root(astrbot_root):
|
||||
raise click.ClickException(
|
||||
|
||||
@@ -5,7 +5,7 @@ from pathlib import Path
|
||||
|
||||
import click
|
||||
|
||||
from ..utils import get_astrbot_root
|
||||
from astrbot.core.utils.astrbot_path import astrbot_paths
|
||||
|
||||
|
||||
@click.command()
|
||||
@@ -51,22 +51,22 @@ def uninstall(yes: bool, keep_data: bool) -> None:
|
||||
)
|
||||
|
||||
# 2. Remove Data
|
||||
astrbot_root = get_astrbot_root()
|
||||
data_dir = astrbot_root / "data"
|
||||
dot_astrbot = astrbot_root / ".astrbot"
|
||||
lock_file = astrbot_root / "astrbot.lock"
|
||||
|
||||
if keep_data:
|
||||
click.echo("Skipping data removal as requested.")
|
||||
return
|
||||
|
||||
# Helper paths
|
||||
dot_astrbot = astrbot_paths.root / ".astrbot"
|
||||
lock_file = astrbot_paths.root / "astrbot.lock"
|
||||
data_dir = astrbot_paths.data
|
||||
|
||||
# Check if this looks like an AstrBot root before blowing things up
|
||||
if not dot_astrbot.exists() and not data_dir.exists():
|
||||
click.echo("No AstrBot initialization found in current directory.")
|
||||
return
|
||||
|
||||
if yes or click.confirm(
|
||||
f"Are you sure you want to remove AstrBot data at {astrbot_root}? \n"
|
||||
f"Are you sure you want to remove AstrBot data at {astrbot_paths.root}? \n"
|
||||
f"This will delete:\n"
|
||||
f" - {data_dir} (Config, Plugins, Database)\n"
|
||||
f" - {dot_astrbot}\n"
|
||||
@@ -87,3 +87,5 @@ def uninstall(yes: bool, keep_data: bool) -> None:
|
||||
lock_file.unlink()
|
||||
|
||||
click.echo("AstrBot data removed successfully.")
|
||||
click.echo("uv: uv tool uninstall astrbot")
|
||||
click.echo("paru/yay: paru -R astrbot")
|
||||
|
||||
@@ -1,9 +1,13 @@
|
||||
from importlib import resources
|
||||
from pathlib import Path
|
||||
|
||||
import click
|
||||
|
||||
from astrbot.core.utils.astrbot_path import astrbot_paths
|
||||
|
||||
# Static assets bundled inside the installed wheel (built by hatch_build.py).
|
||||
_BUNDLED_DIST = Path(__file__).parent.parent.parent / "dashboard" / "dist"
|
||||
# _BUNDLED_DIST = Path(__file__).parent.parent.parent / "dashboard" / "dist"
|
||||
_BUNDLED_DIST = resources.files("astrbot") / "dashboard" / "dist"
|
||||
|
||||
|
||||
def check_astrbot_root(path: str | Path) -> bool:
|
||||
@@ -19,7 +23,7 @@ def check_astrbot_root(path: str | Path) -> bool:
|
||||
|
||||
def get_astrbot_root() -> Path:
|
||||
"""Get the AstrBot root directory path"""
|
||||
return Path.cwd()
|
||||
return astrbot_paths.root
|
||||
|
||||
|
||||
async def check_dashboard(astrbot_root: Path) -> None:
|
||||
@@ -30,7 +34,7 @@ async def check_dashboard(astrbot_root: Path) -> None:
|
||||
from .version_comparator import VersionComparator
|
||||
|
||||
# If the wheel ships bundled dashboard assets, no network download is needed.
|
||||
if _BUNDLED_DIST.exists():
|
||||
if _BUNDLED_DIST.is_dir():
|
||||
click.echo("Dashboard is bundled with the package – skipping download.")
|
||||
return
|
||||
|
||||
|
||||
@@ -14,6 +14,8 @@ Skills 目录路径:固定为数据目录下的 skills 目录
|
||||
"""
|
||||
|
||||
import os
|
||||
from importlib import resources
|
||||
from pathlib import Path
|
||||
|
||||
from astrbot.core.utils.runtime_env import is_packaged_desktop_runtime
|
||||
|
||||
@@ -31,6 +33,7 @@ def get_astrbot_root() -> str:
|
||||
return os.path.realpath(path)
|
||||
if is_packaged_desktop_runtime():
|
||||
return os.path.realpath(os.path.join(os.path.expanduser("~"), ".astrbot"))
|
||||
|
||||
return os.path.realpath(os.getcwd())
|
||||
|
||||
|
||||
@@ -87,3 +90,67 @@ def get_astrbot_knowledge_base_path() -> str:
|
||||
def get_astrbot_backups_path() -> str:
|
||||
"""获取Astrbot备份目录路径"""
|
||||
return os.path.realpath(os.path.join(get_astrbot_data_path(), "backups"))
|
||||
|
||||
|
||||
class AstrbotPaths:
|
||||
"""Astrbot 项目路径管理类"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._root = self._resolve_root()
|
||||
|
||||
def _resolve_root(self) -> Path:
|
||||
if path := os.environ.get("ASTRBOT_ROOT"):
|
||||
return Path(path)
|
||||
if is_packaged_desktop_runtime():
|
||||
return Path().home() / ".astrbot"
|
||||
|
||||
return Path(os.getcwd())
|
||||
|
||||
@property
|
||||
def root(self) -> Path:
|
||||
return self._root
|
||||
|
||||
@root.setter
|
||||
def root(self, value: Path) -> None:
|
||||
self._root = value
|
||||
|
||||
@property
|
||||
def project_root(self) -> Path:
|
||||
"""获取项目根目录路径 (package root)"""
|
||||
with resources.as_file(resources.files("astrbot")) as path:
|
||||
return path
|
||||
|
||||
@property
|
||||
def data(self) -> Path:
|
||||
return self.root / "data"
|
||||
|
||||
@property
|
||||
def config(self) -> Path:
|
||||
return self.data / "config"
|
||||
|
||||
@property
|
||||
def plugins(self) -> Path:
|
||||
return self.data / "plugins"
|
||||
|
||||
@property
|
||||
def temp(self) -> Path:
|
||||
return self.data / "temp"
|
||||
|
||||
@property
|
||||
def skills(self) -> Path:
|
||||
return self.data / "skills"
|
||||
|
||||
@property
|
||||
def site_packages(self) -> Path:
|
||||
return self.data / "site-packages"
|
||||
|
||||
@property
|
||||
def knowledge_base(self) -> Path:
|
||||
return self.data / "knowledge_base"
|
||||
|
||||
@property
|
||||
def backups(self) -> Path:
|
||||
return self.data / "backups"
|
||||
|
||||
|
||||
astrbot_paths = AstrbotPaths()
|
||||
|
||||
213
tests/cli/test_bk.py
Normal file
213
tests/cli/test_bk.py
Normal file
@@ -0,0 +1,213 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
|
||||
from astrbot.cli.commands.cmd_bk import bk
|
||||
from astrbot.core.backup.importer import ImportResult
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_exporter():
|
||||
with patch("astrbot.cli.commands.cmd_bk.AstrBotExporter") as mock:
|
||||
exporter_instance = mock.return_value
|
||||
# export_all is async, return a fake path string
|
||||
exporter_instance.export_all = AsyncMock(return_value="fake_backup.zip")
|
||||
yield exporter_instance
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_importer():
|
||||
with patch("astrbot.cli.commands.cmd_bk.AstrBotImporter") as mock:
|
||||
importer_instance = mock.return_value
|
||||
# import_all is async
|
||||
result = ImportResult()
|
||||
importer_instance.import_all = AsyncMock(return_value=result)
|
||||
yield importer_instance
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_kb_manager():
|
||||
with patch(
|
||||
"astrbot.cli.commands.cmd_bk._get_kb_manager", new_callable=AsyncMock
|
||||
) as mock:
|
||||
mock.return_value = MagicMock()
|
||||
yield mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_gpg_tools():
|
||||
"""Mock shutil.which and asyncio.create_subprocess_exec for GPG operations"""
|
||||
# Create a mock process object
|
||||
mock_process = MagicMock()
|
||||
mock_process.wait = AsyncMock(return_value=None)
|
||||
mock_process.returncode = 0
|
||||
|
||||
with patch("astrbot.cli.commands.cmd_bk.shutil.which") as mock_which, patch(
|
||||
"asyncio.create_subprocess_exec", new_callable=AsyncMock
|
||||
) as mock_exec:
|
||||
mock_which.return_value = "/usr/bin/gpg"
|
||||
mock_exec.return_value = mock_process
|
||||
yield mock_which, mock_exec
|
||||
|
||||
|
||||
def test_export_simple(mock_exporter, mock_kb_manager):
|
||||
"""Test basic export command"""
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(bk, ["export"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "Raw backup exported to: fake_backup.zip" in result.output
|
||||
mock_exporter.export_all.assert_called_once()
|
||||
|
||||
|
||||
def test_export_custom_output(mock_exporter, mock_kb_manager):
|
||||
"""Test export with output directory"""
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(bk, ["export", "-o", "/tmp/backups"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
mock_exporter.export_all.assert_called_once()
|
||||
assert mock_exporter.export_all.call_args[0][0] == "/tmp/backups"
|
||||
|
||||
|
||||
def test_export_gpg_sign(mock_exporter, mock_kb_manager, mock_gpg_tools):
|
||||
"""Test export with GPG signing"""
|
||||
_, mock_exec = mock_gpg_tools
|
||||
runner = CliRunner()
|
||||
|
||||
# Mock Path operations used in GPG block
|
||||
with patch("pathlib.Path.unlink") as mock_unlink, patch(
|
||||
"pathlib.Path.exists", return_value=True
|
||||
):
|
||||
|
||||
result = runner.invoke(bk, ["export", "--gpg-sign"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
mock_exec.assert_called()
|
||||
# create_subprocess_exec(*cmd) -> args are the command parts as separate args
|
||||
args = mock_exec.call_args[0]
|
||||
# Verify GPG command construction
|
||||
assert args[0] == "gpg"
|
||||
assert "--sign" in args
|
||||
assert "--yes" in args
|
||||
# Should clean up original file
|
||||
mock_unlink.assert_called()
|
||||
|
||||
|
||||
def test_export_gpg_encrypt(mock_exporter, mock_kb_manager, mock_gpg_tools):
|
||||
"""Test export with GPG asymmetric encryption"""
|
||||
_, mock_exec = mock_gpg_tools
|
||||
runner = CliRunner()
|
||||
|
||||
with patch("pathlib.Path.unlink"), patch("pathlib.Path.exists", return_value=True):
|
||||
result = runner.invoke(bk, ["export", "--gpg-encrypt", "user@example.com"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
args = mock_exec.call_args[0]
|
||||
assert "--encrypt" in args
|
||||
assert "--recipient" in args
|
||||
assert "user@example.com" in args
|
||||
|
||||
|
||||
def test_export_gpg_symmetric(mock_exporter, mock_kb_manager, mock_gpg_tools):
|
||||
"""Test export with GPG symmetric encryption"""
|
||||
_, mock_exec = mock_gpg_tools
|
||||
runner = CliRunner()
|
||||
|
||||
with patch("pathlib.Path.unlink"), patch("pathlib.Path.exists", return_value=True):
|
||||
result = runner.invoke(bk, ["export", "--gpg-symmetric"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
args = mock_exec.call_args[0]
|
||||
assert "--symmetric" in args
|
||||
assert "--yes" in args
|
||||
|
||||
|
||||
def test_export_digest(mock_exporter, mock_kb_manager):
|
||||
"""Test export with digest generation"""
|
||||
runner = CliRunner()
|
||||
|
||||
# Mock file operations for digest calculation
|
||||
mock_data = b"test data for checksum"
|
||||
with patch("builtins.open", new_callable=MagicMock) as mock_open, patch(
|
||||
"pathlib.Path.write_text"
|
||||
) as mock_write_text:
|
||||
|
||||
# Mock reading file content
|
||||
file_handle = mock_open.return_value.__enter__.return_value
|
||||
file_handle.read.side_effect = [mock_data, b""] # Data then EOF
|
||||
|
||||
result = runner.invoke(bk, ["export", "--digest", "sha256"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "Digest generated" in result.output
|
||||
|
||||
# Verify hash calculation
|
||||
expected_hash = hashlib.sha256(mock_data).hexdigest()
|
||||
mock_write_text.assert_called_once()
|
||||
content = mock_write_text.call_args[0][0]
|
||||
assert expected_hash in content
|
||||
assert "fake_backup.zip" in content
|
||||
|
||||
|
||||
def test_import_simple(mock_importer, mock_kb_manager):
|
||||
"""Test basic import command"""
|
||||
runner = CliRunner()
|
||||
|
||||
with patch("pathlib.Path.exists", return_value=True):
|
||||
result = runner.invoke(bk, ["import", "backup.zip", "--yes"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "Import completed successfully" in result.output
|
||||
mock_importer.import_all.assert_called_once()
|
||||
|
||||
|
||||
def test_import_decrypt(mock_importer, mock_kb_manager, mock_gpg_tools):
|
||||
"""Test import with GPG decryption"""
|
||||
_, mock_exec = mock_gpg_tools
|
||||
runner = CliRunner()
|
||||
|
||||
# path.exists needs to return True for initial check,
|
||||
# then unlink needs to be mocked for cleanup
|
||||
with patch("pathlib.Path.exists", return_value=True), patch(
|
||||
"pathlib.Path.unlink"
|
||||
) as mock_unlink:
|
||||
|
||||
result = runner.invoke(bk, ["import", "backup.zip.gpg", "--yes"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "Processing GPG file" in result.output
|
||||
|
||||
# Verify GPG decryption call
|
||||
mock_exec.assert_called()
|
||||
args = mock_exec.call_args[0]
|
||||
assert args[0] == "gpg"
|
||||
assert "--decrypt" in args
|
||||
assert "backup.zip.gpg" in str(args[-1])
|
||||
|
||||
# Verify temp file cleanup
|
||||
mock_unlink.assert_called()
|
||||
|
||||
|
||||
def test_import_abort(mock_importer):
|
||||
"""Test import abort on confirmation decline"""
|
||||
runner = CliRunner()
|
||||
with patch("pathlib.Path.exists", return_value=True):
|
||||
result = runner.invoke(bk, ["import", "backup.zip"], input="n\n")
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "Aborted" in result.output
|
||||
mock_importer.import_all.assert_not_called()
|
||||
|
||||
|
||||
def test_export_gpg_missing(mock_exporter, mock_kb_manager):
|
||||
"""Test error when GPG is missing"""
|
||||
with patch("astrbot.cli.commands.cmd_bk.shutil.which", return_value=None):
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(bk, ["export", "--gpg-sign"])
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "GPG tool not found" in result.output
|
||||
138
tests/cli/test_uninstall.py
Normal file
138
tests/cli/test_uninstall.py
Normal file
@@ -0,0 +1,138 @@
|
||||
import platform
|
||||
import shutil
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, call, patch
|
||||
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
|
||||
from astrbot.cli.commands.cmd_uninstall import uninstall
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_systemctl():
|
||||
"""Mock shutil.which('systemctl') and subprocess.run"""
|
||||
with patch("astrbot.cli.commands.cmd_uninstall.shutil.which") as mock_which, patch(
|
||||
"astrbot.cli.commands.cmd_uninstall.subprocess.run"
|
||||
) as mock_run:
|
||||
mock_which.return_value = "/usr/bin/systemctl"
|
||||
yield mock_which, mock_run
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_astrbot_paths(tmp_path):
|
||||
"""Mock astrbot_paths to use a temporary directory"""
|
||||
with patch("astrbot.cli.commands.cmd_uninstall.astrbot_paths") as mock_paths:
|
||||
# Create a fake astrbot root structure in tmp_path
|
||||
root = tmp_path / "astrbot_root"
|
||||
root.mkdir()
|
||||
data = root / "data"
|
||||
data.mkdir()
|
||||
dot_astrbot = root / ".astrbot"
|
||||
dot_astrbot.touch()
|
||||
lock_file = root / "astrbot.lock"
|
||||
lock_file.touch()
|
||||
|
||||
mock_paths.root = root
|
||||
mock_paths.data = data
|
||||
yield mock_paths
|
||||
|
||||
|
||||
@pytest.mark.skipif(platform.system() != "Linux", reason="Systemd tests only on Linux")
|
||||
def test_uninstall_systemd_service(mock_systemctl, mock_astrbot_paths):
|
||||
"""Test systemd service removal"""
|
||||
mock_which, mock_run = mock_systemctl
|
||||
runner = CliRunner()
|
||||
|
||||
# Mock Path.home() to return a temp directory
|
||||
with patch("pathlib.Path.home") as mock_home:
|
||||
fake_home = mock_astrbot_paths.root / "home"
|
||||
fake_home.mkdir()
|
||||
mock_home.return_value = fake_home
|
||||
|
||||
# Create fake service file
|
||||
service_dir = fake_home / ".config" / "systemd" / "user"
|
||||
service_dir.mkdir(parents=True)
|
||||
service_file = service_dir / "astrbot.service"
|
||||
service_file.write_text("fake service content")
|
||||
|
||||
# Run with --yes to skip confirmation, --keep-data to focus on systemd
|
||||
result = runner.invoke(uninstall, ["--yes", "--keep-data"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "Stopping AstrBot service..." in result.output
|
||||
assert "Systemd service uninstalled." in result.output
|
||||
|
||||
# Verify subprocess calls
|
||||
mock_run.assert_any_call(
|
||||
["systemctl", "--user", "stop", "astrbot"], check=False
|
||||
)
|
||||
mock_run.assert_any_call(
|
||||
["systemctl", "--user", "disable", "astrbot"], check=False
|
||||
)
|
||||
mock_run.assert_any_call(["systemctl", "--user", "daemon-reload"], check=True)
|
||||
|
||||
# Verify file removed
|
||||
assert not service_file.exists()
|
||||
|
||||
|
||||
def test_uninstall_data_removal(mock_astrbot_paths):
|
||||
"""Test data directory removal"""
|
||||
runner = CliRunner()
|
||||
|
||||
# Verify pre-state
|
||||
assert mock_astrbot_paths.data.exists()
|
||||
assert (mock_astrbot_paths.root / ".astrbot").exists()
|
||||
assert (mock_astrbot_paths.root / "astrbot.lock").exists()
|
||||
|
||||
# Run uninstall with --yes
|
||||
result = runner.invoke(uninstall, ["--yes"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "AstrBot data removed successfully" in result.output
|
||||
|
||||
# Verify removal
|
||||
assert not mock_astrbot_paths.data.exists()
|
||||
assert not (mock_astrbot_paths.root / ".astrbot").exists()
|
||||
assert not (mock_astrbot_paths.root / "astrbot.lock").exists()
|
||||
|
||||
|
||||
def test_uninstall_keep_data(mock_astrbot_paths):
|
||||
"""Test uninstall with --keep-data"""
|
||||
runner = CliRunner()
|
||||
|
||||
result = runner.invoke(uninstall, ["--yes", "--keep-data"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "Skipping data removal as requested" in result.output
|
||||
|
||||
# Verify data still exists
|
||||
assert mock_astrbot_paths.data.exists()
|
||||
|
||||
|
||||
def test_uninstall_abort_on_no_confirm(mock_astrbot_paths):
|
||||
"""Test abort when user declines confirmation"""
|
||||
runner = CliRunner()
|
||||
|
||||
# Input "n" to decline
|
||||
result = runner.invoke(uninstall, input="n\n")
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "Aborted" in result.output
|
||||
|
||||
# Verify data still exists
|
||||
assert mock_astrbot_paths.data.exists()
|
||||
|
||||
|
||||
def test_uninstall_not_astrbot_root(mock_astrbot_paths):
|
||||
"""Test running uninstall in a non-AstrBot directory"""
|
||||
# Remove the marker files
|
||||
shutil.rmtree(mock_astrbot_paths.data)
|
||||
(mock_astrbot_paths.root / ".astrbot").unlink()
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(uninstall, ["--yes"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "No AstrBot initialization found" in result.output
|
||||
Reference in New Issue
Block a user