mirror of
https://github.com/EstrellaXD/Auto_Bangumi.git
synced 2026-11-01 13:15:59 +08:00
feat(downloader): make aria2 a first-class downloader backend
Replace the aria2 client's stub RPC layer with a typed, per-call-timeout JSON-RPC implementation and real add/query/rename/move/delete support: honest add-result tracking with local dedup (url/file-hash) plus RPC duplicate-error detection, torrents_info/torrents_files mapped into the qB-shaped dict the rest of the app expects, and a new aria2_gid table (migration v11) persisting the "ab:<bangumi_id>" tag association so Renamer's tag-based offset lookup works unmodified. Capability flags (can_query/can_rename/can_manage) now reflect what is actually implemented; can_rss_rules stays False (no qB-native RSS-rule surface to honestly back it with). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GTnuEgR78cBk32bLNKar3z
This commit is contained in:
74
backend/src/module/database/aria2.py
Normal file
74
backend/src/module/database/aria2.py
Normal file
@@ -0,0 +1,74 @@
|
||||
"""aria2 gid 关联表数据库操作层。
|
||||
|
||||
只有 ``Aria2Downloader`` 会用到这张表:aria2 没有 tag/category,需要本地
|
||||
持久化 gid -> bangumi_id / category / dedup_key 的映射。
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlmodel import select
|
||||
|
||||
from module.models.aria2 import Aria2Gid
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Aria2GidDatabase:
|
||||
def __init__(self, session: AsyncSession):
|
||||
self.session = session
|
||||
|
||||
async def upsert(
|
||||
self,
|
||||
gid: str,
|
||||
bangumi_id: int | None = None,
|
||||
category: str | None = None,
|
||||
dedup_key: str | None = None,
|
||||
) -> None:
|
||||
"""新增一条 gid 记录,或者用非空字段覆盖已有记录。"""
|
||||
existing = await self.session.get(Aria2Gid, gid)
|
||||
if existing is None:
|
||||
self.session.add(
|
||||
Aria2Gid(
|
||||
gid=gid,
|
||||
bangumi_id=bangumi_id,
|
||||
category=category,
|
||||
dedup_key=dedup_key,
|
||||
)
|
||||
)
|
||||
else:
|
||||
if bangumi_id is not None:
|
||||
existing.bangumi_id = bangumi_id
|
||||
if category is not None:
|
||||
existing.category = category
|
||||
if dedup_key is not None:
|
||||
existing.dedup_key = dedup_key
|
||||
self.session.add(existing)
|
||||
await self.session.commit()
|
||||
|
||||
async def get(self, gid: str) -> Aria2Gid | None:
|
||||
return await self.session.get(Aria2Gid, gid)
|
||||
|
||||
async def get_many(self, gids: list[str]) -> dict[str, Aria2Gid]:
|
||||
if not gids:
|
||||
return {}
|
||||
result = await self.session.execute(
|
||||
select(Aria2Gid).where(Aria2Gid.gid.in_(gids)) # type: ignore[attr-defined]
|
||||
)
|
||||
return {row.gid: row for row in result.scalars().all()}
|
||||
|
||||
async def find_by_dedup_key(self, dedup_key: str) -> str | None:
|
||||
"""返回携带该 dedup_key 的已有 gid(没有则 None),用于新增前判重。"""
|
||||
result = await self.session.execute(
|
||||
select(Aria2Gid.gid).where(Aria2Gid.dedup_key == dedup_key) # type: ignore[arg-type]
|
||||
)
|
||||
return result.scalars().first()
|
||||
|
||||
async def set_category(self, gid: str, category: str) -> None:
|
||||
await self.upsert(gid, category=category)
|
||||
|
||||
async def delete(self, gid: str) -> None:
|
||||
existing = await self.session.get(Aria2Gid, gid)
|
||||
if existing is not None:
|
||||
await self.session.delete(existing)
|
||||
await self.session.commit()
|
||||
@@ -6,6 +6,7 @@ from sqlmodel import SQLModel
|
||||
|
||||
from module.models import Bangumi, User
|
||||
|
||||
from .aria2 import Aria2GidDatabase
|
||||
from .bangumi import BangumiDatabase
|
||||
from .engine import async_engine, async_session_factory
|
||||
from .migrations import ( # noqa: F401 (re-exported for existing importers)
|
||||
@@ -41,6 +42,7 @@ class Database:
|
||||
self.torrent = TorrentDatabase(self.session)
|
||||
self.bangumi = BangumiDatabase(self.session)
|
||||
self.user = UserDatabase(self.session)
|
||||
self.aria2 = Aria2GidDatabase(self.session)
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
@@ -155,6 +155,21 @@ MIGRATIONS: tuple[Migration, ...] = (
|
||||
),
|
||||
column_exists("bangumi", "preferred_group"),
|
||||
),
|
||||
Migration(
|
||||
11,
|
||||
"create aria2_gid table for aria2 gid<->bangumi association",
|
||||
(
|
||||
"""CREATE TABLE IF NOT EXISTS aria2_gid (
|
||||
gid VARCHAR NOT NULL PRIMARY KEY,
|
||||
bangumi_id INTEGER REFERENCES bangumi(id),
|
||||
category VARCHAR,
|
||||
dedup_key VARCHAR,
|
||||
created_at TIMESTAMP NOT NULL
|
||||
)""",
|
||||
"CREATE INDEX IF NOT EXISTS ix_aria2_gid_dedup_key ON aria2_gid(dedup_key)",
|
||||
),
|
||||
table_exists("aria2_gid"),
|
||||
),
|
||||
)
|
||||
|
||||
# 由迁移列表派生,新增迁移时无需手动同步
|
||||
|
||||
@@ -1,22 +1,52 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
from typing import ClassVar
|
||||
|
||||
import httpx
|
||||
|
||||
from module.database import Database
|
||||
from module.downloader.base import DownloaderCapabilities
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 每次 JSON-RPC 调用的默认超时(秒),auth() 期间的探测调用会传入更短的值。
|
||||
_DEFAULT_CALL_TIMEOUT = 10.0
|
||||
|
||||
# qBittorrent 风格 status_filter -> aria2 status 集合的映射。None(不过滤)
|
||||
# 之外的未知取值一律当作"不过滤"处理,而不是返回空列表。
|
||||
_STATUS_FILTER_MAP: dict[str, set[str]] = {
|
||||
"completed": {"complete"},
|
||||
"downloading": {"active"},
|
||||
"active": {"active"},
|
||||
"paused": {"paused"},
|
||||
"errored": {"error"},
|
||||
}
|
||||
|
||||
|
||||
class Aria2RpcError(Exception):
|
||||
"""aria2 JSON-RPC 返回的业务错误(服务器收到了请求,但拒绝执行)。"""
|
||||
|
||||
def __init__(self, code: int | None, message: str):
|
||||
self.code = code
|
||||
self.message = message
|
||||
super().__init__(f"aria2 RPC error {code}: {message}")
|
||||
|
||||
|
||||
class Aria2ConnectionError(Exception):
|
||||
"""请求根本没有得到 aria2 服务器的有效响应(网络/超时/JSON 解析失败)。"""
|
||||
|
||||
|
||||
class Aria2Downloader:
|
||||
# Aria2 only speaks "add a download": it has no torrent query, rename,
|
||||
# management, or qB-native RSS surface. The facade consults this and skips
|
||||
# everything else instead of calling methods that do not exist here.
|
||||
# aria2 通过 JSON-RPC 暴露真实的查询/重命名/管理能力(tellActive/tellFiles/
|
||||
# changeOption 等),只是没有 qBittorrent 原生的 RSS 规则功能,因此
|
||||
# can_rss_rules 保持 False,其余按真实支持情况置 True。
|
||||
capabilities: ClassVar[DownloaderCapabilities] = DownloaderCapabilities(
|
||||
can_query=False,
|
||||
can_rename=False,
|
||||
can_manage=False,
|
||||
can_query=True,
|
||||
can_rename=True,
|
||||
can_manage=True,
|
||||
can_rss_rules=False,
|
||||
)
|
||||
|
||||
@@ -28,26 +58,41 @@ class Aria2Downloader:
|
||||
self._rpc_url = f"{host}/jsonrpc"
|
||||
self._id = 0
|
||||
|
||||
async def _call(self, method: str, params: list | None = None):
|
||||
async def _call(
|
||||
self,
|
||||
method: str,
|
||||
params: list | None = None,
|
||||
timeout: float = _DEFAULT_CALL_TIMEOUT,
|
||||
):
|
||||
assert self._client is not None, "Aria2Downloader.auth() must run first"
|
||||
self._id += 1
|
||||
if params is None:
|
||||
params = []
|
||||
# Prepend token
|
||||
full_params = [f"token:{self.secret}"] + params
|
||||
full_params = [f"token:{self.secret}"] + (params or [])
|
||||
payload = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": self._id,
|
||||
"method": f"aria2.{method}",
|
||||
"params": full_params,
|
||||
}
|
||||
resp = await self._client.post(self._rpc_url, json=payload)
|
||||
result = resp.json()
|
||||
try:
|
||||
resp = await self._client.post(self._rpc_url, json=payload, timeout=timeout)
|
||||
except httpx.TimeoutException as e:
|
||||
raise Aria2ConnectionError(f"aria2 RPC '{method}' timed out: {e}") from e
|
||||
except httpx.RequestError as e:
|
||||
raise Aria2ConnectionError(
|
||||
f"aria2 RPC '{method}' request failed: {e}"
|
||||
) from e
|
||||
try:
|
||||
result = resp.json()
|
||||
except ValueError as e:
|
||||
raise Aria2ConnectionError(
|
||||
f"aria2 RPC '{method}' returned invalid JSON: {e}"
|
||||
) from e
|
||||
if "error" in result:
|
||||
raise Exception(f"Aria2 RPC error: {result['error']}")
|
||||
err = result["error"] or {}
|
||||
raise Aria2RpcError(err.get("code"), err.get("message", str(err)))
|
||||
return result.get("result")
|
||||
|
||||
async def auth(self, retry=3):
|
||||
async def auth(self, retry: int = 3) -> bool:
|
||||
if self._client is not None and self._authed:
|
||||
return True
|
||||
if self._client is None:
|
||||
@@ -60,7 +105,7 @@ class Aria2Downloader:
|
||||
await self._call("getVersion")
|
||||
self._authed = True
|
||||
return True
|
||||
except Exception as e:
|
||||
except (Aria2RpcError, Aria2ConnectionError) as e:
|
||||
logger.warning(
|
||||
f"Can't login Aria2 Server {self.host}, retry in 5 seconds. Error: {e}"
|
||||
)
|
||||
@@ -68,30 +113,389 @@ class Aria2Downloader:
|
||||
times += 1
|
||||
return False
|
||||
|
||||
async def logout(self):
|
||||
async def logout(self) -> None:
|
||||
self._authed = False
|
||||
if self._client:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
async def torrents_files(self, torrent_hash: str):
|
||||
return []
|
||||
async def check_connection(self) -> str:
|
||||
version = await self._call("getVersion")
|
||||
return (version or {}).get("version", "unknown")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _normalize_hashes(hashes: str | list[str] | tuple[str, ...]) -> list[str]:
|
||||
if isinstance(hashes, (list, tuple)):
|
||||
return list(hashes)
|
||||
return hashes.split("|") if "|" in hashes else [hashes]
|
||||
|
||||
@staticmethod
|
||||
def _parse_bangumi_id_from_tag(tags: str | None) -> int | None:
|
||||
"""从 'ab:<id>' 格式的 tag/tags 中解出 bangumi_id(与 Renamer 的约定一致)。"""
|
||||
if not tags:
|
||||
return None
|
||||
for tag in tags.split(","):
|
||||
tag = tag.strip()
|
||||
if tag.startswith("ab:"):
|
||||
try:
|
||||
return int(tag[3:])
|
||||
except ValueError:
|
||||
pass
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _is_duplicate_error(e: Aria2RpcError) -> bool:
|
||||
msg = e.message.lower()
|
||||
return "already" in msg or "duplicate" in msg
|
||||
|
||||
@staticmethod
|
||||
def _is_not_found_error(e: Aria2RpcError) -> bool:
|
||||
msg = e.message.lower()
|
||||
return "not found" in msg or "is not found" in msg
|
||||
|
||||
@staticmethod
|
||||
def _extract_name(download: dict) -> str:
|
||||
bt_info = (download.get("bittorrent") or {}).get("info") or {}
|
||||
if bt_info.get("name"):
|
||||
return bt_info["name"]
|
||||
files = download.get("files") or []
|
||||
if files and files[0].get("path"):
|
||||
return os.path.basename(files[0]["path"])
|
||||
return download.get("gid", "")
|
||||
|
||||
@staticmethod
|
||||
def _move_file(old_abs: str, new_abs: str) -> None:
|
||||
os.makedirs(os.path.dirname(new_abs), exist_ok=True)
|
||||
os.replace(old_abs, new_abs)
|
||||
|
||||
@staticmethod
|
||||
def _move_files(files: list[dict], old_dir: str, new_dir: str) -> None:
|
||||
old_dir = os.path.normpath(old_dir)
|
||||
for f in files:
|
||||
path = f.get("path")
|
||||
if not path or not os.path.exists(path):
|
||||
continue
|
||||
rel = os.path.relpath(path, old_dir)
|
||||
dest = os.path.join(new_dir, rel)
|
||||
os.makedirs(os.path.dirname(dest), exist_ok=True)
|
||||
shutil.move(path, dest)
|
||||
|
||||
@staticmethod
|
||||
def _remove_file_and_empty_dirs(path: str, boundary: str) -> None:
|
||||
"""删除单个文件,并向上清理变空的父目录,直到 boundary(含)为止不再删除。
|
||||
|
||||
boundary 通常是 aria2 的 ``dir``(bangumi 的 season 目录),可能被多个
|
||||
gid 共享,所以绝不删除 boundary 本身或它之外的目录。
|
||||
"""
|
||||
if os.path.exists(path):
|
||||
os.remove(path)
|
||||
boundary = os.path.normpath(boundary)
|
||||
parent = os.path.dirname(path)
|
||||
while os.path.normpath(parent) != boundary and (
|
||||
os.path.normpath(parent).startswith(boundary + os.sep)
|
||||
):
|
||||
try:
|
||||
os.rmdir(parent)
|
||||
except OSError:
|
||||
break
|
||||
parent = os.path.dirname(parent)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Adding
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _add_uri(self, url: str, options: dict) -> str | None:
|
||||
try:
|
||||
return await self._call("addUri", [[url], options])
|
||||
except Aria2RpcError as e:
|
||||
if self._is_duplicate_error(e):
|
||||
logger.debug("[Aria2] addUri reports duplicate: %s", e.message)
|
||||
else:
|
||||
logger.error("[Aria2] addUri failed for %s: %s", url, e)
|
||||
return None
|
||||
except Aria2ConnectionError as e:
|
||||
logger.error("[Aria2] addUri connection error for %s: %s", url, e)
|
||||
return None
|
||||
|
||||
async def _add_torrent_file(self, data: bytes, options: dict) -> str | None:
|
||||
import base64
|
||||
|
||||
b64 = base64.b64encode(data).decode()
|
||||
try:
|
||||
return await self._call("addTorrent", [b64, [], options])
|
||||
except Aria2RpcError as e:
|
||||
if self._is_duplicate_error(e):
|
||||
logger.debug("[Aria2] addTorrent reports duplicate: %s", e.message)
|
||||
else:
|
||||
logger.error("[Aria2] addTorrent failed: %s", e)
|
||||
return None
|
||||
except Aria2ConnectionError as e:
|
||||
logger.error("[Aria2] addTorrent connection error: %s", e)
|
||||
return None
|
||||
|
||||
async def add_torrents(
|
||||
self, torrent_urls, torrent_files, save_path, category, tags=None
|
||||
):
|
||||
import base64
|
||||
) -> bool:
|
||||
"""添加下载任务,返回 True 当且仅当真的新增了至少一个任务。
|
||||
|
||||
去重优先靠本地的 dedup_key(url 或种子内容 hash)判断"已经添加过",
|
||||
aria2 RPC 报错里带 already/duplicate 字样时也当作重复处理(见
|
||||
``_is_duplicate_error``)——这样调用方(DownloadClient.add_torrent)
|
||||
才能像 qBittorrent 一样区分"新增成功"和"之前加过了"。
|
||||
"""
|
||||
bangumi_id = self._parse_bangumi_id_from_tag(tags)
|
||||
options = {"dir": save_path}
|
||||
if torrent_urls:
|
||||
urls = torrent_urls if isinstance(torrent_urls, list) else [torrent_urls]
|
||||
for url in urls:
|
||||
await self._call("addUri", [[url], options])
|
||||
if torrent_files:
|
||||
files = (
|
||||
torrent_files if isinstance(torrent_files, list) else [torrent_files]
|
||||
added_any = False
|
||||
async with Database() as db:
|
||||
if torrent_urls:
|
||||
urls = (
|
||||
torrent_urls if isinstance(torrent_urls, list) else [torrent_urls]
|
||||
)
|
||||
for url in urls:
|
||||
dedup_key = f"url:{url}"
|
||||
if await db.aria2.find_by_dedup_key(dedup_key):
|
||||
logger.debug("[Aria2] Skip already-added url: %s", url)
|
||||
continue
|
||||
gid = await self._add_uri(url, options)
|
||||
if gid is None:
|
||||
continue
|
||||
await db.aria2.upsert(gid, bangumi_id, category, dedup_key)
|
||||
added_any = True
|
||||
if torrent_files:
|
||||
files = (
|
||||
torrent_files
|
||||
if isinstance(torrent_files, list)
|
||||
else [torrent_files]
|
||||
)
|
||||
for f in files:
|
||||
dedup_key = f"file:{hashlib.sha1(f).hexdigest()}"
|
||||
if await db.aria2.find_by_dedup_key(dedup_key):
|
||||
logger.debug("[Aria2] Skip already-added torrent file")
|
||||
continue
|
||||
gid = await self._add_torrent_file(f, options)
|
||||
if gid is None:
|
||||
continue
|
||||
await db.aria2.upsert(gid, bangumi_id, category, dedup_key)
|
||||
added_any = True
|
||||
return added_any
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Querying
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def torrents_info(self, status_filter, category, tag=None) -> list[dict]:
|
||||
try:
|
||||
active = await self._call("tellActive") or []
|
||||
waiting = await self._call("tellWaiting", [0, 1000]) or []
|
||||
stopped = await self._call("tellStopped", [0, 1000]) or []
|
||||
except (Aria2RpcError, Aria2ConnectionError) as e:
|
||||
logger.error("[Aria2] Failed to query downloads: %s", e)
|
||||
return []
|
||||
raw = [*active, *waiting, *stopped]
|
||||
gids = [d["gid"] for d in raw if d.get("gid")]
|
||||
async with Database() as db:
|
||||
meta = await db.aria2.get_many(gids)
|
||||
|
||||
allowed_statuses = (
|
||||
_STATUS_FILTER_MAP.get(status_filter) if status_filter else None
|
||||
)
|
||||
result = []
|
||||
for d in raw:
|
||||
gid = d.get("gid", "")
|
||||
aria2_status = d.get("status", "")
|
||||
if allowed_statuses is not None and aria2_status not in allowed_statuses:
|
||||
continue
|
||||
info = meta.get(gid)
|
||||
torrent_category = info.category if info else None
|
||||
if category and torrent_category != category:
|
||||
continue
|
||||
tags_str = f"ab:{info.bangumi_id}" if info and info.bangumi_id else ""
|
||||
if tag and tag != tags_str:
|
||||
continue
|
||||
result.append(
|
||||
{
|
||||
"hash": gid,
|
||||
"name": self._extract_name(d),
|
||||
"save_path": d.get("dir", ""),
|
||||
"tags": tags_str,
|
||||
"category": torrent_category or "",
|
||||
"state": aria2_status,
|
||||
"size": int(d.get("totalLength", 0) or 0),
|
||||
"progress": (
|
||||
int(d.get("completedLength", 0) or 0) / int(d["totalLength"])
|
||||
if d.get("totalLength")
|
||||
else 0
|
||||
),
|
||||
"dlspeed": int(d.get("downloadSpeed", 0) or 0),
|
||||
"upspeed": int(d.get("uploadSpeed", 0) or 0),
|
||||
}
|
||||
)
|
||||
for f in files:
|
||||
b64 = base64.b64encode(f).decode()
|
||||
await self._call("addTorrent", [b64, [], options])
|
||||
return result
|
||||
|
||||
async def torrents_files(self, torrent_hash: str) -> list[dict]:
|
||||
try:
|
||||
files = await self._call("getFiles", [torrent_hash])
|
||||
except (Aria2RpcError, Aria2ConnectionError) as e:
|
||||
logger.error("[Aria2] Failed to get files for %s: %s", torrent_hash, e)
|
||||
return []
|
||||
save_dir = ""
|
||||
try:
|
||||
status = await self._call("tellStatus", [torrent_hash, ["dir"]])
|
||||
save_dir = (status or {}).get("dir", "")
|
||||
except (Aria2RpcError, Aria2ConnectionError) as e:
|
||||
logger.debug("[Aria2] Could not resolve dir for %s: %s", torrent_hash, e)
|
||||
result = []
|
||||
for f in files or []:
|
||||
path = f.get("path", "")
|
||||
rel = os.path.relpath(path, save_dir) if save_dir and path else path
|
||||
result.append({"name": rel, "size": int(f.get("length", 0) or 0)})
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Rename (filesystem move -- aria2 downloads land in a known local dir)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def torrents_rename_file(
|
||||
self, torrent_hash, old_path, new_path, verify: bool = True
|
||||
) -> bool:
|
||||
try:
|
||||
status = await self._call("tellStatus", [torrent_hash, ["dir"]])
|
||||
except (Aria2RpcError, Aria2ConnectionError) as e:
|
||||
logger.warning(
|
||||
"[Aria2] Rename failed, cannot resolve save dir for %s: %s",
|
||||
torrent_hash,
|
||||
e,
|
||||
)
|
||||
return False
|
||||
save_dir = (status or {}).get("dir", "")
|
||||
if not save_dir:
|
||||
return False
|
||||
old_abs = os.path.join(save_dir, old_path)
|
||||
new_abs = os.path.join(save_dir, new_path)
|
||||
try:
|
||||
await asyncio.to_thread(self._move_file, old_abs, new_abs)
|
||||
except OSError as e:
|
||||
logger.warning(
|
||||
"[Aria2] Failed to rename file %s -> %s: %s", old_abs, new_abs, e
|
||||
)
|
||||
return False
|
||||
if verify:
|
||||
exists = await asyncio.to_thread(os.path.exists, new_abs)
|
||||
if not exists:
|
||||
logger.debug(
|
||||
"[Aria2] Rename reported success but %s is missing", new_abs
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Management
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def torrents_delete(self, hash, delete_files: bool = True) -> bool:
|
||||
gids = self._normalize_hashes(hash)
|
||||
ok = True
|
||||
async with Database() as db:
|
||||
for gid in gids:
|
||||
save_dir = None
|
||||
try:
|
||||
status = await self._call("tellStatus", [gid, ["dir"]])
|
||||
save_dir = (status or {}).get("dir")
|
||||
except (Aria2RpcError, Aria2ConnectionError) as e:
|
||||
logger.debug("[Aria2] Could not resolve dir for %s: %s", gid, e)
|
||||
|
||||
if delete_files and save_dir:
|
||||
try:
|
||||
files = await self._call("getFiles", [gid])
|
||||
except (Aria2RpcError, Aria2ConnectionError):
|
||||
files = []
|
||||
for f in files or []:
|
||||
path = f.get("path")
|
||||
if not path:
|
||||
continue
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
self._remove_file_and_empty_dirs, path, save_dir
|
||||
)
|
||||
except OSError as e:
|
||||
logger.warning(
|
||||
"[Aria2] Failed to delete file %s: %s", path, e
|
||||
)
|
||||
|
||||
try:
|
||||
await self._call("forceRemove", [gid])
|
||||
except Aria2RpcError as e:
|
||||
if not self._is_not_found_error(e):
|
||||
logger.error("[Aria2] Failed to remove %s: %s", gid, e)
|
||||
ok = False
|
||||
except Aria2ConnectionError as e:
|
||||
logger.error("[Aria2] Failed to remove %s: %s", gid, e)
|
||||
ok = False
|
||||
|
||||
await db.aria2.delete(gid)
|
||||
return ok
|
||||
|
||||
async def torrents_pause(self, hashes: str | list[str]) -> None:
|
||||
for gid in self._normalize_hashes(hashes):
|
||||
try:
|
||||
await self._call("forcePause", [gid])
|
||||
except (Aria2RpcError, Aria2ConnectionError) as e:
|
||||
logger.warning("[Aria2] Failed to pause %s: %s", gid, e)
|
||||
|
||||
async def torrents_resume(self, hashes: str | list[str]) -> None:
|
||||
for gid in self._normalize_hashes(hashes):
|
||||
try:
|
||||
await self._call("unpause", [gid])
|
||||
except (Aria2RpcError, Aria2ConnectionError) as e:
|
||||
logger.warning("[Aria2] Failed to resume %s: %s", gid, e)
|
||||
|
||||
async def move_torrent(self, hashes, new_location) -> None:
|
||||
for gid in self._normalize_hashes(hashes):
|
||||
try:
|
||||
status = await self._call("tellStatus", [gid, ["dir"]])
|
||||
except (Aria2RpcError, Aria2ConnectionError) as e:
|
||||
logger.warning(
|
||||
"[Aria2] Cannot move %s, status lookup failed: %s", gid, e
|
||||
)
|
||||
continue
|
||||
old_dir = (status or {}).get("dir")
|
||||
if not old_dir or os.path.normpath(old_dir) == os.path.normpath(
|
||||
new_location
|
||||
):
|
||||
continue
|
||||
try:
|
||||
files = await self._call("getFiles", [gid])
|
||||
except (Aria2RpcError, Aria2ConnectionError) as e:
|
||||
logger.warning("[Aria2] Cannot move %s, file list failed: %s", gid, e)
|
||||
continue
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
self._move_files, files or [], old_dir, new_location
|
||||
)
|
||||
except OSError as e:
|
||||
logger.warning("[Aria2] Failed to move files for %s: %s", gid, e)
|
||||
continue
|
||||
try:
|
||||
await self._call("changeOption", [gid, {"dir": new_location}])
|
||||
except (Aria2RpcError, Aria2ConnectionError) as e:
|
||||
logger.debug(
|
||||
"[Aria2] changeOption dir failed for %s (non-fatal): %s", gid, e
|
||||
)
|
||||
|
||||
async def set_category(self, _hash, category) -> None:
|
||||
async with Database() as db:
|
||||
for gid in self._normalize_hashes(_hash):
|
||||
await db.aria2.set_category(gid, category)
|
||||
|
||||
async def add_tag(self, _hash: str, tag: str) -> None:
|
||||
"""只认识 'ab:<bangumi_id>' 格式的 tag(用于 offset 关联),其余忽略。"""
|
||||
bangumi_id = self._parse_bangumi_id_from_tag(tag)
|
||||
if bangumi_id is None:
|
||||
logger.debug("[Aria2] Ignoring unsupported tag: %s", tag)
|
||||
return
|
||||
async with Database() as db:
|
||||
await db.aria2.upsert(_hash, bangumi_id=bangumi_id)
|
||||
|
||||
@@ -117,10 +117,11 @@ class DownloadClient:
|
||||
elif downloader_type == "aria2":
|
||||
from .client.aria2_downloader import Aria2Downloader
|
||||
|
||||
# Aria2Downloader only implements the capability-gated subset it
|
||||
# declares via `capabilities` (all False) -- the facade never
|
||||
# calls the rest, but that makes it structurally narrower than
|
||||
# the full `DownloaderClient` protocol.
|
||||
# Aria2Downloader implements query/rename/manage for real (see its
|
||||
# `capabilities`), but has no qB-native RSS-rule/prefs surface
|
||||
# (can_rss_rules=False), so it stays structurally narrower than
|
||||
# the full `DownloaderClient` protocol -- the facade skips the
|
||||
# rss/prefs methods it never calls on this backend.
|
||||
return Aria2Downloader(host, username, password) # type: ignore[return-value]
|
||||
elif downloader_type == "mock":
|
||||
from .client.mock_downloader import MockDownloader
|
||||
|
||||
23
backend/src/module/models/aria2.py
Normal file
23
backend/src/module/models/aria2.py
Normal file
@@ -0,0 +1,23 @@
|
||||
"""aria2 gid <-> bangumi 关联表模型。
|
||||
|
||||
aria2 没有 qBittorrent 的 tag/category 概念,只能通过本地记录把下载任务的
|
||||
gid 关联回 bangumi_id(用于续传后的 offset 查找)以及伪造的 category(用于
|
||||
torrents_info 的分类过滤)和 dedup_key(用于识别"已经添加过")。
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Optional
|
||||
|
||||
from sqlmodel import Field, SQLModel
|
||||
|
||||
|
||||
class Aria2Gid(SQLModel, table=True):
|
||||
"""aria2 下载任务 (gid) 到 bangumi / category / 去重键 的本地映射。"""
|
||||
|
||||
__tablename__ = "aria2_gid"
|
||||
|
||||
gid: str = Field(primary_key=True)
|
||||
bangumi_id: Optional[int] = Field(default=None, foreign_key="bangumi.id")
|
||||
category: Optional[str] = None
|
||||
dedup_key: Optional[str] = Field(default=None, index=True)
|
||||
created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
764
backend/src/test/test_aria2_downloader.py
Normal file
764
backend/src/test/test_aria2_downloader.py
Normal file
@@ -0,0 +1,764 @@
|
||||
"""Tests for Aria2Downloader: JSON-RPC error handling, add/query/rename/manage.
|
||||
|
||||
Mirrors test_qb_downloader.py's approach: low-level RPC behaviour is verified
|
||||
by patching ``self._client`` directly (mock JSON-RPC transport); higher-level
|
||||
business logic (dedup, gid<->bangumi association, filesystem rename/move) is
|
||||
verified by patching ``Aria2Downloader._call`` and letting the real logic
|
||||
(including the real, per-test temp-file DB from conftest's autouse
|
||||
``_bind_bare_database`` fixture) run.
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from module.database import Database
|
||||
from module.downloader.client.aria2_downloader import (
|
||||
Aria2ConnectionError,
|
||||
Aria2Downloader,
|
||||
Aria2RpcError,
|
||||
)
|
||||
|
||||
|
||||
def _aria2() -> Aria2Downloader:
|
||||
return Aria2Downloader("http://localhost:6800", "u", "secret-token")
|
||||
|
||||
|
||||
def _rpc_response(result=None, error=None):
|
||||
resp = MagicMock()
|
||||
body: dict = {"jsonrpc": "2.0", "id": 1}
|
||||
if error is not None:
|
||||
body["error"] = error
|
||||
else:
|
||||
body["result"] = result
|
||||
resp.json.return_value = body
|
||||
return resp
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _call: low-level JSON-RPC transport
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCallRpcLayer:
|
||||
async def test_call_success_returns_result(self):
|
||||
aria2 = _aria2()
|
||||
aria2._client = AsyncMock()
|
||||
aria2._client.post = AsyncMock(return_value=_rpc_response(result="1.36.0"))
|
||||
|
||||
result = await aria2._call("getVersion")
|
||||
|
||||
assert result == "1.36.0"
|
||||
|
||||
async def test_call_prepends_token_as_first_param(self):
|
||||
aria2 = _aria2()
|
||||
aria2._client = AsyncMock()
|
||||
aria2._client.post = AsyncMock(return_value=_rpc_response(result=[]))
|
||||
|
||||
await aria2._call("tellActive", [1, 2])
|
||||
|
||||
payload = aria2._client.post.call_args.kwargs["json"]
|
||||
assert payload["params"] == ["token:secret-token", 1, 2]
|
||||
assert payload["method"] == "aria2.tellActive"
|
||||
|
||||
async def test_call_error_response_raises_aria2_rpc_error(self):
|
||||
aria2 = _aria2()
|
||||
aria2._client = AsyncMock()
|
||||
aria2._client.post = AsyncMock(
|
||||
return_value=_rpc_response(error={"code": 1, "message": "Unauthorized"})
|
||||
)
|
||||
|
||||
with pytest.raises(Aria2RpcError) as excinfo:
|
||||
await aria2._call("getVersion")
|
||||
assert excinfo.value.code == 1
|
||||
assert excinfo.value.message == "Unauthorized"
|
||||
|
||||
async def test_call_timeout_raises_connection_error(self):
|
||||
aria2 = _aria2()
|
||||
aria2._client = AsyncMock()
|
||||
aria2._client.post = AsyncMock(side_effect=httpx.TimeoutException("timed out"))
|
||||
|
||||
with pytest.raises(Aria2ConnectionError):
|
||||
await aria2._call("getVersion")
|
||||
|
||||
async def test_call_request_error_raises_connection_error(self):
|
||||
aria2 = _aria2()
|
||||
aria2._client = AsyncMock()
|
||||
aria2._client.post = AsyncMock(side_effect=httpx.ConnectError("refused"))
|
||||
|
||||
with pytest.raises(Aria2ConnectionError):
|
||||
await aria2._call("getVersion")
|
||||
|
||||
async def test_call_invalid_json_raises_connection_error(self):
|
||||
aria2 = _aria2()
|
||||
aria2._client = AsyncMock()
|
||||
bad_resp = MagicMock()
|
||||
bad_resp.json.side_effect = ValueError("not json")
|
||||
aria2._client.post = AsyncMock(return_value=bad_resp)
|
||||
|
||||
with pytest.raises(Aria2ConnectionError):
|
||||
await aria2._call("getVersion")
|
||||
|
||||
async def test_call_uses_default_timeout_when_not_overridden(self):
|
||||
aria2 = _aria2()
|
||||
aria2._client = AsyncMock()
|
||||
aria2._client.post = AsyncMock(return_value=_rpc_response(result=[]))
|
||||
|
||||
await aria2._call("tellActive")
|
||||
|
||||
assert aria2._client.post.call_args.kwargs["timeout"] == 10.0
|
||||
|
||||
async def test_call_honors_explicit_timeout_override(self):
|
||||
aria2 = _aria2()
|
||||
aria2._client = AsyncMock()
|
||||
aria2._client.post = AsyncMock(return_value=_rpc_response(result=[]))
|
||||
|
||||
await aria2._call("getVersion", timeout=2.5)
|
||||
|
||||
assert aria2._client.post.call_args.kwargs["timeout"] == 2.5
|
||||
|
||||
async def test_call_without_auth_raises_assertion(self):
|
||||
aria2 = _aria2()
|
||||
with pytest.raises(AssertionError):
|
||||
await aria2._call("getVersion")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# auth()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAuth:
|
||||
async def test_auth_returns_true_on_success(self):
|
||||
aria2 = _aria2()
|
||||
with patch.object(aria2, "_call", AsyncMock(return_value={"version": "1.0"})):
|
||||
result = await aria2.auth()
|
||||
assert result is True
|
||||
assert aria2._authed is True
|
||||
|
||||
async def test_auth_reuses_already_authed_client(self):
|
||||
aria2 = _aria2()
|
||||
call_mock = AsyncMock(return_value={"version": "1.0"})
|
||||
with patch.object(aria2, "_call", call_mock):
|
||||
await aria2.auth()
|
||||
await aria2.auth()
|
||||
assert call_mock.call_count == 1
|
||||
|
||||
async def test_auth_retries_and_returns_false_on_persistent_rpc_error(self):
|
||||
aria2 = _aria2()
|
||||
with (
|
||||
patch.object(
|
||||
aria2, "_call", AsyncMock(side_effect=Aria2RpcError(1, "bad token"))
|
||||
),
|
||||
patch(
|
||||
"module.downloader.client.aria2_downloader.asyncio.sleep",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
):
|
||||
result = await aria2.auth(retry=2)
|
||||
assert result is False
|
||||
assert aria2._authed is False
|
||||
|
||||
async def test_auth_retries_on_connection_error_then_succeeds(self):
|
||||
aria2 = _aria2()
|
||||
call_mock = AsyncMock(
|
||||
side_effect=[Aria2ConnectionError("refused"), {"version": "1.0"}]
|
||||
)
|
||||
with (
|
||||
patch.object(aria2, "_call", call_mock),
|
||||
patch(
|
||||
"module.downloader.client.aria2_downloader.asyncio.sleep",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
):
|
||||
result = await aria2.auth(retry=3)
|
||||
assert result is True
|
||||
assert call_mock.call_count == 2
|
||||
|
||||
async def test_logout_clears_client_and_authed_flag(self):
|
||||
aria2 = _aria2()
|
||||
with patch.object(aria2, "_call", AsyncMock(return_value={"version": "1.0"})):
|
||||
await aria2.auth()
|
||||
await aria2.logout()
|
||||
assert aria2._client is None
|
||||
assert aria2._authed is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# add_torrents: honest results + gid tracking + dedup
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAddTorrents:
|
||||
async def test_new_url_returns_true_and_persists_gid(self):
|
||||
aria2 = _aria2()
|
||||
with patch.object(aria2, "_call", AsyncMock(return_value="gid001")):
|
||||
result = await aria2.add_torrents(
|
||||
torrent_urls="magnet:?xt=urn:btih:abc",
|
||||
torrent_files=None,
|
||||
save_path="/downloads/Show/Season 1",
|
||||
category="Bangumi",
|
||||
tags="ab:42",
|
||||
)
|
||||
assert result is True
|
||||
async with Database() as db:
|
||||
record = await db.aria2.get("gid001")
|
||||
assert record is not None
|
||||
assert record.bangumi_id == 42
|
||||
assert record.category == "Bangumi"
|
||||
|
||||
async def test_duplicate_url_second_call_returns_false_and_skips_rpc(self):
|
||||
aria2 = _aria2()
|
||||
call_mock = AsyncMock(return_value="gid001")
|
||||
with patch.object(aria2, "_call", call_mock):
|
||||
await aria2.add_torrents(
|
||||
torrent_urls="magnet:?xt=urn:btih:abc",
|
||||
torrent_files=None,
|
||||
save_path="/downloads",
|
||||
category="Bangumi",
|
||||
tags="ab:1",
|
||||
)
|
||||
result = await aria2.add_torrents(
|
||||
torrent_urls="magnet:?xt=urn:btih:abc",
|
||||
torrent_files=None,
|
||||
save_path="/downloads",
|
||||
category="Bangumi",
|
||||
tags="ab:1",
|
||||
)
|
||||
assert result is False
|
||||
# addUri only called once -- the second, duplicate add never hit aria2.
|
||||
assert call_mock.call_count == 1
|
||||
|
||||
async def test_partial_duplicate_returns_true(self):
|
||||
aria2 = _aria2()
|
||||
call_mock = AsyncMock(side_effect=["gid001", "gid002"])
|
||||
with patch.object(aria2, "_call", call_mock):
|
||||
await aria2.add_torrents(
|
||||
torrent_urls=["magnet:?xt=urn:btih:aaa"],
|
||||
torrent_files=None,
|
||||
save_path="/downloads",
|
||||
category="Bangumi",
|
||||
tags="ab:1",
|
||||
)
|
||||
result = await aria2.add_torrents(
|
||||
torrent_urls=["magnet:?xt=urn:btih:aaa", "magnet:?xt=urn:btih:bbb"],
|
||||
torrent_files=None,
|
||||
save_path="/downloads",
|
||||
category="Bangumi",
|
||||
tags="ab:1",
|
||||
)
|
||||
assert result is True
|
||||
assert call_mock.call_count == 2
|
||||
|
||||
async def test_all_urls_already_added_returns_false(self):
|
||||
aria2 = _aria2()
|
||||
call_mock = AsyncMock(return_value="gid001")
|
||||
with patch.object(aria2, "_call", call_mock):
|
||||
await aria2.add_torrents(
|
||||
torrent_urls=["magnet:?xt=urn:btih:aaa"],
|
||||
torrent_files=None,
|
||||
save_path="/downloads",
|
||||
category="Bangumi",
|
||||
tags="ab:1",
|
||||
)
|
||||
result = await aria2.add_torrents(
|
||||
torrent_urls=["magnet:?xt=urn:btih:aaa"],
|
||||
torrent_files=None,
|
||||
save_path="/downloads",
|
||||
category="Bangumi",
|
||||
tags="ab:1",
|
||||
)
|
||||
assert result is False
|
||||
|
||||
async def test_rpc_duplicate_error_treated_as_already_added(self):
|
||||
aria2 = _aria2()
|
||||
with patch.object(
|
||||
aria2,
|
||||
"_call",
|
||||
AsyncMock(side_effect=Aria2RpcError(1, "GID already exists")),
|
||||
):
|
||||
result = await aria2.add_torrents(
|
||||
torrent_urls="magnet:?xt=urn:btih:zzz",
|
||||
torrent_files=None,
|
||||
save_path="/downloads",
|
||||
category="Bangumi",
|
||||
)
|
||||
assert result is False
|
||||
async with Database() as db:
|
||||
# Nothing to persist -- addUri never returned a gid.
|
||||
assert (
|
||||
await db.aria2.find_by_dedup_key("url:magnet:?xt=urn:btih:zzz") is None
|
||||
)
|
||||
|
||||
async def test_rpc_non_duplicate_error_returns_false_for_that_item(self):
|
||||
aria2 = _aria2()
|
||||
with patch.object(
|
||||
aria2, "_call", AsyncMock(side_effect=Aria2RpcError(3, "Disk full"))
|
||||
):
|
||||
result = await aria2.add_torrents(
|
||||
torrent_urls="magnet:?xt=urn:btih:zzz",
|
||||
torrent_files=None,
|
||||
save_path="/downloads",
|
||||
category="Bangumi",
|
||||
)
|
||||
assert result is False
|
||||
|
||||
async def test_torrent_file_bytes_dedup_by_content_hash(self):
|
||||
aria2 = _aria2()
|
||||
call_mock = AsyncMock(return_value="gid-file-1")
|
||||
payload = b"fake torrent bytes"
|
||||
with patch.object(aria2, "_call", call_mock):
|
||||
first = await aria2.add_torrents(
|
||||
torrent_urls=None,
|
||||
torrent_files=payload,
|
||||
save_path="/downloads",
|
||||
category="Bangumi",
|
||||
tags="ab:7",
|
||||
)
|
||||
second = await aria2.add_torrents(
|
||||
torrent_urls=None,
|
||||
torrent_files=payload,
|
||||
save_path="/downloads",
|
||||
category="Bangumi",
|
||||
tags="ab:7",
|
||||
)
|
||||
assert first is True
|
||||
assert second is False
|
||||
assert call_mock.call_count == 1
|
||||
|
||||
async def test_no_tags_leaves_bangumi_id_unset(self):
|
||||
aria2 = _aria2()
|
||||
with patch.object(aria2, "_call", AsyncMock(return_value="gid-no-tag")):
|
||||
await aria2.add_torrents(
|
||||
torrent_urls="magnet:?xt=urn:btih:notag",
|
||||
torrent_files=None,
|
||||
save_path="/downloads",
|
||||
category="Bangumi",
|
||||
tags=None,
|
||||
)
|
||||
async with Database() as db:
|
||||
record = await db.aria2.get("gid-no-tag")
|
||||
assert record is not None
|
||||
assert record.bangumi_id is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# torrents_info: aria2 -> qB-shaped dict mapping
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _download(gid, status="complete", dir_="/downloads/Show/Season 1", name=None):
|
||||
d = {
|
||||
"gid": gid,
|
||||
"status": status,
|
||||
"dir": dir_,
|
||||
"totalLength": "1000",
|
||||
"completedLength": "1000" if status == "complete" else "500",
|
||||
"downloadSpeed": "0",
|
||||
"uploadSpeed": "0",
|
||||
"files": [{"path": f"{dir_}/{name or gid}.mkv"}],
|
||||
}
|
||||
return d
|
||||
|
||||
|
||||
class TestTorrentsInfo:
|
||||
async def test_maps_fields_to_qb_shape(self):
|
||||
aria2 = _aria2()
|
||||
|
||||
async def fake_call(method, params=None, timeout=10.0):
|
||||
if method == "tellActive":
|
||||
return [_download("gidA", status="active")]
|
||||
return []
|
||||
|
||||
with patch.object(aria2, "_call", AsyncMock(side_effect=fake_call)):
|
||||
async with Database() as db:
|
||||
await db.aria2.upsert("gidA", bangumi_id=5, category="Bangumi")
|
||||
result = await aria2.torrents_info(status_filter=None, category=None)
|
||||
|
||||
assert len(result) == 1
|
||||
info = result[0]
|
||||
assert info["hash"] == "gidA"
|
||||
assert info["save_path"] == "/downloads/Show/Season 1"
|
||||
assert info["tags"] == "ab:5"
|
||||
assert info["category"] == "Bangumi"
|
||||
assert info["state"] == "active"
|
||||
assert info["name"] == "gidA.mkv"
|
||||
|
||||
async def test_status_filter_completed_excludes_active(self):
|
||||
aria2 = _aria2()
|
||||
|
||||
async def fake_call(method, params=None, timeout=10.0):
|
||||
if method == "tellActive":
|
||||
return [_download("gidActive", status="active")]
|
||||
if method == "tellStopped":
|
||||
return [_download("gidDone", status="complete")]
|
||||
return []
|
||||
|
||||
with patch.object(aria2, "_call", AsyncMock(side_effect=fake_call)):
|
||||
result = await aria2.torrents_info(status_filter="completed", category=None)
|
||||
|
||||
assert [r["hash"] for r in result] == ["gidDone"]
|
||||
|
||||
async def test_status_filter_none_includes_everything(self):
|
||||
aria2 = _aria2()
|
||||
|
||||
async def fake_call(method, params=None, timeout=10.0):
|
||||
if method == "tellActive":
|
||||
return [_download("gidActive", status="active")]
|
||||
if method == "tellStopped":
|
||||
return [_download("gidDone", status="complete")]
|
||||
return []
|
||||
|
||||
with patch.object(aria2, "_call", AsyncMock(side_effect=fake_call)):
|
||||
result = await aria2.torrents_info(status_filter=None, category=None)
|
||||
|
||||
assert {r["hash"] for r in result} == {"gidActive", "gidDone"}
|
||||
|
||||
async def test_category_filter_excludes_other_categories(self):
|
||||
aria2 = _aria2()
|
||||
|
||||
async def fake_call(method, params=None, timeout=10.0):
|
||||
if method == "tellActive":
|
||||
return [_download("gidA"), _download("gidB")]
|
||||
return []
|
||||
|
||||
with patch.object(aria2, "_call", AsyncMock(side_effect=fake_call)):
|
||||
async with Database() as db:
|
||||
await db.aria2.upsert("gidA", category="Bangumi")
|
||||
await db.aria2.upsert("gidB", category="Other")
|
||||
result = await aria2.torrents_info(status_filter=None, category="Bangumi")
|
||||
|
||||
assert [r["hash"] for r in result] == ["gidA"]
|
||||
|
||||
async def test_untracked_gid_has_no_tag_and_empty_category(self):
|
||||
aria2 = _aria2()
|
||||
|
||||
async def fake_call(method, params=None, timeout=10.0):
|
||||
if method == "tellActive":
|
||||
return [_download("gidUntracked")]
|
||||
return []
|
||||
|
||||
with patch.object(aria2, "_call", AsyncMock(side_effect=fake_call)):
|
||||
result = await aria2.torrents_info(status_filter=None, category=None)
|
||||
|
||||
assert result[0]["tags"] == ""
|
||||
assert result[0]["category"] == ""
|
||||
|
||||
async def test_query_failure_returns_empty_list(self):
|
||||
aria2 = _aria2()
|
||||
with patch.object(
|
||||
aria2, "_call", AsyncMock(side_effect=Aria2ConnectionError("down"))
|
||||
):
|
||||
result = await aria2.torrents_info(status_filter=None, category=None)
|
||||
assert result == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# torrents_files: relative-path mapping
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTorrentsFiles:
|
||||
async def test_returns_paths_relative_to_dir(self):
|
||||
aria2 = _aria2()
|
||||
|
||||
async def fake_call(method, params=None, timeout=10.0):
|
||||
if method == "getFiles":
|
||||
return [
|
||||
{"path": "/downloads/Show/Season 1/ep01.mkv", "length": "100"},
|
||||
{"path": "/downloads/Show/Season 1/ep01.ass", "length": "1"},
|
||||
]
|
||||
if method == "tellStatus":
|
||||
return {"dir": "/downloads/Show/Season 1"}
|
||||
return None
|
||||
|
||||
with patch.object(aria2, "_call", AsyncMock(side_effect=fake_call)):
|
||||
files = await aria2.torrents_files("gidA")
|
||||
|
||||
assert files == [
|
||||
{"name": "ep01.mkv", "size": 100},
|
||||
{"name": "ep01.ass", "size": 1},
|
||||
]
|
||||
|
||||
async def test_query_failure_returns_empty_list(self):
|
||||
aria2 = _aria2()
|
||||
with patch.object(
|
||||
aria2, "_call", AsyncMock(side_effect=Aria2ConnectionError("down"))
|
||||
):
|
||||
files = await aria2.torrents_files("gidA")
|
||||
assert files == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# torrents_rename_file: real filesystem move
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTorrentsRenameFile:
|
||||
async def test_renames_file_on_disk(self, tmp_path):
|
||||
aria2 = _aria2()
|
||||
(tmp_path / "old.mkv").write_bytes(b"data")
|
||||
|
||||
async def fake_call(method, params=None, timeout=10.0):
|
||||
return {"dir": str(tmp_path)}
|
||||
|
||||
with patch.object(aria2, "_call", AsyncMock(side_effect=fake_call)):
|
||||
result = await aria2.torrents_rename_file("gidA", "old.mkv", "new.mkv")
|
||||
|
||||
assert result is True
|
||||
assert not (tmp_path / "old.mkv").exists()
|
||||
assert (tmp_path / "new.mkv").read_bytes() == b"data"
|
||||
|
||||
async def test_creates_parent_dirs_for_nested_new_path(self, tmp_path):
|
||||
aria2 = _aria2()
|
||||
(tmp_path / "old.mkv").write_bytes(b"data")
|
||||
|
||||
async def fake_call(method, params=None, timeout=10.0):
|
||||
return {"dir": str(tmp_path)}
|
||||
|
||||
with patch.object(aria2, "_call", AsyncMock(side_effect=fake_call)):
|
||||
result = await aria2.torrents_rename_file(
|
||||
"gidA", "old.mkv", os.path.join("Season 01", "new.mkv")
|
||||
)
|
||||
|
||||
assert result is True
|
||||
assert (tmp_path / "Season 01" / "new.mkv").exists()
|
||||
|
||||
async def test_returns_false_when_source_file_missing(self, tmp_path):
|
||||
aria2 = _aria2()
|
||||
|
||||
async def fake_call(method, params=None, timeout=10.0):
|
||||
return {"dir": str(tmp_path)}
|
||||
|
||||
with patch.object(aria2, "_call", AsyncMock(side_effect=fake_call)):
|
||||
result = await aria2.torrents_rename_file("gidA", "missing.mkv", "new.mkv")
|
||||
|
||||
assert result is False
|
||||
|
||||
async def test_returns_false_when_status_lookup_fails(self):
|
||||
aria2 = _aria2()
|
||||
with patch.object(
|
||||
aria2, "_call", AsyncMock(side_effect=Aria2ConnectionError("down"))
|
||||
):
|
||||
result = await aria2.torrents_rename_file("gidA", "old.mkv", "new.mkv")
|
||||
assert result is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# torrents_delete: filesystem cleanup that never touches sibling gids
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTorrentsDelete:
|
||||
async def test_deletes_only_this_gids_files_and_prunes_empty_subdir(self, tmp_path):
|
||||
aria2 = _aria2()
|
||||
season_dir = tmp_path / "Season 1"
|
||||
this_gid_dir = season_dir / "MyBatch"
|
||||
this_gid_dir.mkdir(parents=True)
|
||||
this_file = this_gid_dir / "ep01.mkv"
|
||||
this_file.write_bytes(b"data")
|
||||
sibling_file = season_dir / "ep02.mkv"
|
||||
sibling_file.write_bytes(b"other data")
|
||||
|
||||
async def fake_call(method, params=None, timeout=10.0):
|
||||
if method == "tellStatus":
|
||||
return {"dir": str(season_dir)}
|
||||
if method == "getFiles":
|
||||
return [{"path": str(this_file)}]
|
||||
if method == "forceRemove":
|
||||
return "gidA"
|
||||
return None
|
||||
|
||||
async with Database() as db:
|
||||
await db.aria2.upsert("gidA", bangumi_id=1)
|
||||
|
||||
with patch.object(aria2, "_call", AsyncMock(side_effect=fake_call)):
|
||||
result = await aria2.torrents_delete("gidA", delete_files=True)
|
||||
|
||||
assert result is True
|
||||
assert not this_file.exists()
|
||||
assert not this_gid_dir.exists() # emptied subdirectory pruned
|
||||
assert sibling_file.exists() # sibling never touched
|
||||
assert season_dir.exists() # boundary directory never removed
|
||||
|
||||
async with Database() as db:
|
||||
assert await db.aria2.get("gidA") is None
|
||||
|
||||
async def test_returns_false_on_remove_rpc_error(self):
|
||||
aria2 = _aria2()
|
||||
|
||||
async def fake_call(method, params=None, timeout=10.0):
|
||||
if method == "tellStatus":
|
||||
return {"dir": "/downloads"}
|
||||
if method == "getFiles":
|
||||
return []
|
||||
if method == "forceRemove":
|
||||
raise Aria2RpcError(1, "Internal error")
|
||||
return None
|
||||
|
||||
with patch.object(aria2, "_call", AsyncMock(side_effect=fake_call)):
|
||||
result = await aria2.torrents_delete("gidA", delete_files=False)
|
||||
|
||||
assert result is False
|
||||
|
||||
async def test_not_found_error_is_treated_as_already_gone(self):
|
||||
aria2 = _aria2()
|
||||
|
||||
async def fake_call(method, params=None, timeout=10.0):
|
||||
if method == "tellStatus":
|
||||
return {"dir": "/downloads"}
|
||||
if method == "getFiles":
|
||||
return []
|
||||
if method == "forceRemove":
|
||||
raise Aria2RpcError(1, "GID#gidA is not found")
|
||||
return None
|
||||
|
||||
with patch.object(aria2, "_call", AsyncMock(side_effect=fake_call)):
|
||||
result = await aria2.torrents_delete("gidA", delete_files=False)
|
||||
|
||||
assert result is True
|
||||
|
||||
async def test_accepts_pipe_joined_hashes(self):
|
||||
aria2 = _aria2()
|
||||
seen_gids = []
|
||||
|
||||
async def fake_call(method, params=None, timeout=10.0):
|
||||
if method == "tellStatus":
|
||||
return {"dir": "/downloads"}
|
||||
if method == "getFiles":
|
||||
return []
|
||||
if method == "forceRemove":
|
||||
seen_gids.append(params[0])
|
||||
return params[0]
|
||||
return None
|
||||
|
||||
with patch.object(aria2, "_call", AsyncMock(side_effect=fake_call)):
|
||||
await aria2.torrents_delete("gidA|gidB", delete_files=False)
|
||||
|
||||
assert seen_gids == ["gidA", "gidB"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# pause / resume
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPauseResume:
|
||||
async def test_pause_calls_force_pause_per_gid(self):
|
||||
aria2 = _aria2()
|
||||
call_mock = AsyncMock(return_value=None)
|
||||
with patch.object(aria2, "_call", call_mock):
|
||||
await aria2.torrents_pause(["gidA", "gidB"])
|
||||
methods = [c.args[0] for c in call_mock.call_args_list]
|
||||
assert methods == ["forcePause", "forcePause"]
|
||||
|
||||
async def test_resume_calls_unpause_per_gid(self):
|
||||
aria2 = _aria2()
|
||||
call_mock = AsyncMock(return_value=None)
|
||||
with patch.object(aria2, "_call", call_mock):
|
||||
await aria2.torrents_resume("gidA|gidB")
|
||||
methods = [c.args[0] for c in call_mock.call_args_list]
|
||||
assert methods == ["unpause", "unpause"]
|
||||
|
||||
async def test_pause_failure_is_logged_not_raised(self):
|
||||
aria2 = _aria2()
|
||||
with patch.object(
|
||||
aria2, "_call", AsyncMock(side_effect=Aria2ConnectionError("down"))
|
||||
):
|
||||
await aria2.torrents_pause("gidA") # must not raise
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# move_torrent
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMoveTorrent:
|
||||
async def test_moves_files_to_new_dir_and_updates_option(self, tmp_path):
|
||||
aria2 = _aria2()
|
||||
old_dir = tmp_path / "old"
|
||||
new_dir = tmp_path / "new"
|
||||
old_dir.mkdir()
|
||||
file_path = old_dir / "ep01.mkv"
|
||||
file_path.write_bytes(b"data")
|
||||
|
||||
change_option_calls = []
|
||||
|
||||
async def fake_call(method, params=None, timeout=10.0):
|
||||
if method == "tellStatus":
|
||||
return {"dir": str(old_dir)}
|
||||
if method == "getFiles":
|
||||
return [{"path": str(file_path)}]
|
||||
if method == "changeOption":
|
||||
change_option_calls.append(params)
|
||||
return "OK"
|
||||
return None
|
||||
|
||||
with patch.object(aria2, "_call", AsyncMock(side_effect=fake_call)):
|
||||
await aria2.move_torrent("gidA", str(new_dir))
|
||||
|
||||
assert not file_path.exists()
|
||||
assert (new_dir / "ep01.mkv").read_bytes() == b"data"
|
||||
assert change_option_calls == [["gidA", {"dir": str(new_dir)}]]
|
||||
|
||||
async def test_skips_when_new_location_equals_current_dir(self, tmp_path):
|
||||
aria2 = _aria2()
|
||||
call_mock = AsyncMock(return_value={"dir": str(tmp_path)})
|
||||
with patch.object(aria2, "_call", call_mock):
|
||||
await aria2.move_torrent("gidA", str(tmp_path))
|
||||
# Only the tellStatus lookup ran -- no getFiles/changeOption follow-up.
|
||||
assert call_mock.call_args_list[-1].args[0] == "tellStatus"
|
||||
assert call_mock.call_count == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# set_category / add_tag
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSetCategoryAddTag:
|
||||
async def test_set_category_persists_to_sidecar(self):
|
||||
aria2 = _aria2()
|
||||
await aria2.set_category("gidA", "BangumiCollection")
|
||||
async with Database() as db:
|
||||
record = await db.aria2.get("gidA")
|
||||
assert record is not None
|
||||
assert record.category == "BangumiCollection"
|
||||
|
||||
async def test_add_tag_parses_ab_prefix_and_persists_bangumi_id(self):
|
||||
aria2 = _aria2()
|
||||
await aria2.add_tag("gidA", "ab:99")
|
||||
async with Database() as db:
|
||||
record = await db.aria2.get("gidA")
|
||||
assert record is not None
|
||||
assert record.bangumi_id == 99
|
||||
|
||||
async def test_add_tag_ignores_unsupported_tag(self):
|
||||
aria2 = _aria2()
|
||||
await aria2.add_tag("gidA", "some-other-tag")
|
||||
async with Database() as db:
|
||||
record = await db.aria2.get("gidA")
|
||||
assert record is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# check_connection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCheckConnection:
|
||||
async def test_returns_version_string(self):
|
||||
aria2 = _aria2()
|
||||
with patch.object(
|
||||
aria2, "_call", AsyncMock(return_value={"version": "1.36.0"})
|
||||
):
|
||||
result = await aria2.check_connection()
|
||||
assert result == "1.36.0"
|
||||
|
||||
async def test_missing_version_field_returns_unknown(self):
|
||||
aria2 = _aria2()
|
||||
with patch.object(aria2, "_call", AsyncMock(return_value={})):
|
||||
result = await aria2.check_connection()
|
||||
assert result == "unknown"
|
||||
@@ -1,8 +1,8 @@
|
||||
"""Conformance tests for the download-client protocol and capabilities.
|
||||
|
||||
QbDownloader and MockDownloader implement the full async surface; Aria2Downloader
|
||||
only implements the core subset (authenticate + add torrents). All three advertise
|
||||
their capabilities.
|
||||
implements query/rename/manage for real but has no qB-native RSS-rule/prefs
|
||||
surface (can_rss_rules=False). All three advertise their capabilities.
|
||||
"""
|
||||
|
||||
from module.downloader.base import (
|
||||
@@ -38,9 +38,9 @@ class TestCapabilities:
|
||||
can_query=True, can_rename=True, can_manage=True, can_rss_rules=True
|
||||
)
|
||||
|
||||
def test_aria2_supports_only_core(self):
|
||||
def test_aria2_supports_query_rename_manage_but_not_rss_rules(self):
|
||||
assert Aria2Downloader.capabilities == DownloaderCapabilities(
|
||||
can_query=False, can_rename=False, can_manage=False, can_rss_rules=False
|
||||
can_query=True, can_rename=True, can_manage=True, can_rss_rules=False
|
||||
)
|
||||
|
||||
def test_all_three_expose_capabilities(self):
|
||||
|
||||
Reference in New Issue
Block a user