From 2bc0272a512e232c9bc48762c88e2a54bfcfcbe7 Mon Sep 17 00:00:00 2001 From: Maverick Charmer Date: Fri, 31 Oct 2025 03:54:46 +0800 Subject: [PATCH 1/6] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E4=BA=86=E5=AF=B9API?= =?UTF-8?q?=E7=9A=84=E6=B5=8B=E8=AF=95;=E4=BF=AE=E6=94=B9=E4=BA=86WS?= =?UTF-8?q?=E7=9A=84=E7=AE=A1=E7=90=86=E7=B1=BB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 2 +- main.py | 114 +++++++++++++++----------------------- tests/test_api.py | 80 ++++++++++++++++++++++++++ tests/test_jmcomic_lib.py | 5 +- 4 files changed, 129 insertions(+), 72 deletions(-) create mode 100644 tests/test_api.py diff --git a/README.md b/README.md index 87e6309..ff252c0 100644 --- a/README.md +++ b/README.md @@ -57,7 +57,7 @@ RUN mkdir -p /app/temp RUN chmod 777 /app/temp -RUN pip install --no-cache-dir -r requirements.txt +RUN pip install -r requirements.txt EXPOSE 7860 diff --git a/main.py b/main.py index 146a146..ac41e59 100644 --- a/main.py +++ b/main.py @@ -4,12 +4,13 @@ import shutil import asyncio from http.client import responses -from typing import Dict +from typing import Dict, Optional from fastapi import * from starlette.concurrency import run_in_threadpool import uvicorn import jmcomic from pathlib import Path +from fastapi.responses import JSONResponse # --- 全局配置和初始化 --- app = FastAPI() @@ -35,37 +36,32 @@ # --- WebSocket 连接管理器 --- class ConnectionManager: - """管理活跃的 WebSocket 连接，使用 client_id 作为 key""" - + """ + 管理 WebSocket 连接，记录主事件循环并提供线程安全的发送接口。 + """ def __init__(self): - # 存储 client_id 到 WebSocket 对象的映射 self.active_connections: Dict[str, WebSocket] = {} + self.loop: Optional[asyncio.AbstractEventLoop] = None async def connect(self, client_id: str, websocket: WebSocket): - """接受新连接，并将其关联到唯一的 client_id""" await websocket.accept() self.active_connections[client_id] = websocket - print(f"[WebSocket] 新客户端连接 ID: {client_id}") - - def disconnect(self, client_id: str): - """移除断开的连接""" - if client_id in self.active_connections: - del self.active_connections[client_id] - print(f"[WebSocket] 客户端断开连接 ID: {client_id}") - - async def send_personal_message(self, client_id: str, message: dict): - """向特定客户端发送 JSON 消息""" - if client_id in self.active_connections: - websocket = self.active_connections[client_id] + if self.loop is None: + self.loop = asyncio.get_event_loop() + print(f"[WebSocket] 客户端 {client_id} 已连接。") + + async def _send_and_close(self, client_id: str, message: dict): + websocket = self.active_connections.get(client_id) + if websocket: + await websocket.send_json(message) + print(f"[WebSocket] 发送消息给客户端 {client_id}: {message}") try: - await websocket.send_json(message) - print(f"[WebSocket] 成功向客户端 {client_id} 发送通知。") - except Exception as e: - print(f"[WebSocket Error] 向客户端 {client_id} 发送消息失败: {e}") - self.disconnect(client_id) + await websocket.close() + except Exception: + pass + self.active_connections.pop(client_id, None) else: - print(f"[WebSocket Error] 客户端 {client_id} 的连接不存在。") - + print(f"[WebSocket] 未找到客户端 {client_id} 的连接，无法发送消息。") manager = ConnectionManager() @@ -98,15 +94,6 @@ async def websocket_endpoint(websocket: WebSocket, client_id: str): 客户端连接时必须提供唯一的 client_id。 """ await manager.connect(client_id, websocket) - try: - # 保持连接活跃，等待断开 - while True: - await websocket.receive_text() - except WebSocketDisconnect: - manager.disconnect(client_id) - except Exception as e: - print(f"[WebSocket Error] 接收循环出错: {e}") - manager.disconnect(client_id) # --- 阻塞任务处理函数 (在新线程中运行) --- @@ -119,7 +106,6 @@ def sync_download_and_zip_task(album_id: int, client_id: str): print(f"[Task] 开始执行相册 {album_id} 的阻塞下载任务...") try: - # 配置 Jmcomic 选项 (使用 Path 对象来构建 base_dir) optionStr = f""" client: cache: null @@ -157,8 +143,6 @@ def sync_download_and_zip_task(album_id: int, client_id: str): """ option = jmcomic.create_option_by_str(optionStr) jmcomic.JmModuleConfig.CLASS_DOWNLOADER = jmcomic.JmDownloader - - # 阻塞调用 album_list = jmcomic.download_album(album_id, option) if not album_list: @@ -169,64 +153,55 @@ def sync_download_and_zip_task(album_id: int, client_id: str): zip_file_path = FILE_PATH / zip_file_name if zip_file_path.exists(): - # 启动延迟删除线程 (0.5 * 60 * 60 秒 = 30 分钟) - threading.Thread(target=delayed_delete, - args=(zip_file_path, int(0.5 * 60 * 60)), - daemon=True).start() - - # 任务成功，使用 asyncio.run() 在当前线程中运行异步通知任务 - asyncio.run(manager.send_personal_message(client_id, { + message = { "status": "download_ready", "file_name": file_title, "message": f"文件 '{file_title}' 已完成处理，可以下载。" - })) - + } else: - # 任务失败，发送 WebSocket 错误通知 - asyncio.run(manager.send_personal_message(client_id, { + message = { "status": "error", "file_name": file_title, "message": f"文件 '{file_title}' 未找到或处理失败。" - })) - + } + if manager.loop: + future = asyncio.run_coroutine_threadsafe(manager._send_and_close(client_id, message), manager.loop) + try: + future.result(timeout=10) + except Exception as e: + print(f"[Task] 通过主循环发送消息失败: {e}") + else: + print("[Task] 未记录到主事件循环，无法发送 WebSocket 通知。") except Exception as e: - print(f"[Task Error] 下载任务发生异常: {e}") - # 任务异常，发送 WebSocket 错误通知 - asyncio.run(manager.send_personal_message(client_id, { - "status": "error", - "file_name": "", - "message": f"下载任务失败: {str(e)}" - })) + if manager.loop: + fut = asyncio.run_coroutine_threadsafe( + manager._send_and_close(client_id, + {"status": "error", "file_name": "", "message": f"下载任务失败: {str(e)}"}), + manager.loop + ) + try: + fut.result(timeout=10) + except Exception as ee: + print(f"[Task] 发送异常通知失败: {ee}") # --- HTTP 任务启动路由 (替换原 download_album) --- @app.post("/v1/download/album/{album_id}") async def start_album_download(album_id: int, request: Request): - """ - 接收下载请求，立即返回 202，并在后台线程中启动耗时的下载任务。 - 客户端必须在 POST body 中提供唯一的 client_id 来接收通知。 - """ try: data = await request.json() client_id = data.get("client_id") except Exception: raise HTTPException(status_code=400, detail="Request body must be valid JSON containing 'client_id'.") - if not client_id or client_id not in manager.active_connections: - raise HTTPException(status_code=400, detail="WebSocket connection not established or client_id invalid.") - - # 立即在线程池中启动同步阻塞任务 print(f"[Server] 接收下载请求，相册 ID: {album_id}，客户端 ID: {client_id}。任务将在后台启动...") - # 使用 asyncio.create_task 包裹 run_in_threadpool，使其完全在后台异步运行 asyncio.create_task(run_in_threadpool(sync_download_and_zip_task, album_id, client_id)) # 返回 HTTP 202 Accepted 响应，告知客户端任务已接收 - return Response( - status_code=202, - content={"status": "processing", "message": "下载任务已在后台启动，请通过 WebSocket 监听 'download_ready' 通知。"} - ) + return JSONResponse(status_code=202, content={"status": "processing", + "message": "下载任务已在后台启动，请通过 WebSocket 监听 'download_ready' 通知。"}) # --- HTTP 文件下载路由 --- @@ -240,7 +215,6 @@ async def download_file(file_name: str): file_path = FILE_PATH / zip_file_name # 使用统一的 Path 对象和变量 if file_path.exists(): - # 使用 FastAPI.responses.FileResponse return responses.FileResponse(file_path, filename=zip_file_name, media_type="application/zip") return Response( diff --git a/tests/test_api.py b/tests/test_api.py new file mode 100644 index 0000000..c07f8b9 --- /dev/null +++ b/tests/test_api.py @@ -0,0 +1,80 @@ +from fastapi.testclient import TestClient +import time +import os +from pathlib import Path +from main import app + +current_dir = os.getcwd() +FILE_PATH = Path(f"{current_dir}/temp") +os.makedirs(FILE_PATH, exist_ok=True) + +def test_read_root(): + client = TestClient(app) + nowtimestamp = int(time.time() * 1000) + response = client.get("/v1/{0}".format(nowtimestamp)) + timedelta = int(time.time() * 1000) - nowtimestamp + ms = int(timedelta) + assert response.status_code == 200 + assert response.json().get("status") == "ok" + assert response.json().get("app") == "jmcomic_server_api" + assert int(response.json().get("latency")) <= ms + assert int(response.json().get("latency")) > 0 + + +def test_search_album(): + client = TestClient(app) + tag = "全彩" + num = 1 + response = client.get("/v1/search/{0}/{1}".format(tag, num)) + assert response.status_code == 200 + assert isinstance(response.json(), list) + assert len(response.json()) > 0 + first_album = response.json()[0] + assert "album_id" in first_album + assert "title" in first_album + + +def test_get_cover_and_info(): + client = TestClient(app) + aid = 1225432 + response = client.get("/v1/info/{0}".format(aid)) + assert response.status_code == 200 + info_json = response.json() + assert info_json.get("status") == "success" + assert "全彩" in info_json.get("tag", []) + assert int(info_json.get("view_count")) > 0 + assert int(info_json.get("like_count")) > 0 + assert info_json.get("page_count") == "0" + response = client.get("/v1/get/cover/{0}".format(aid)) + assert response.status_code == 200 + assert response.headers["content-type"] == "image/jpeg" + file_path = FILE_PATH / f"cover-{aid}.jpg" + assert file_path.exists() + file_path.unlink() + assert file_path.exists() == False + + +def test_download_album(): + client = TestClient(app) + aid = 1225432 + client_id = "1145141919810" + with client.websocket_connect(f"/ws/notifications/{client_id}") as websocket: + response = client.post(f"/v1/download/album/{aid}", json={"client_id": client_id}) + assert response.status_code == 202 + assert response.json() == { + "status": "processing", + "message": "下载任务已在后台启动，请通过 WebSocket 监听 'download_ready' 通知。" + } + data = websocket.receive_json() + assert data == { + "status": "download_ready", + "file_name": "［酸菜鱼ゅ°］ヒルチャールに败北した胡桃 表情、台词差分", + "message": f"文件 '［酸菜鱼ゅ°］ヒルチャールに败北した胡桃 表情、台词差分' 已完成处理，可以下载。" + } + client.get("/v1/download/album/{0}".format(aid)) + file_title = "［酸菜鱼ゅ°］ヒルチャールに败北した胡桃 表情、台词差分" + zip_file_name = f"{file_title}.zip" + zip_file_path = FILE_PATH / zip_file_name + assert zip_file_path.exists() == True + zip_file_path.unlink() + assert zip_file_path.exists() == False diff --git a/tests/test_jmcomic_lib.py b/tests/test_jmcomic_lib.py index add1822..2eebb07 100644 --- a/tests/test_jmcomic_lib.py +++ b/tests/test_jmcomic_lib.py @@ -6,15 +6,17 @@ FILE_PATH = Path(f"{current_dir}/temp") os.makedirs(FILE_PATH, exist_ok=True) + def test_get_comic_info(): testClient = jmcomic.JmOption.default().new_jm_client() page = testClient.search_site(search_query="1225432") album: jmcomic.JmAlbumDetail = page.single_album assert album.title == "［酸菜鱼ゅ°］ヒルチャールに败北した胡桃 表情、台词差分" - assert album.tags == ["全彩","贫乳","调教","中文"] + assert album.tags == ["全彩", "贫乳", "调教", "中文"] assert album.views is not None assert album.likes is not None + def test_rank_comic(): client = jmcomic.JmOption.default().new_jm_client() page1: jmcomic.JmCategoryPage = client.month_ranking(1) @@ -24,6 +26,7 @@ def test_rank_comic(): assert page2.page_size > 0 assert page3.page_size > 0 + def test_comic_download(): optionStr = f""" client: From 2a676481cf673b20b7e531f544cae4f7bea539dc Mon Sep 17 00:00:00 2001 From: Maverick Charmer Date: Fri, 31 Oct 2025 03:58:13 +0800 Subject: [PATCH 2/6] =?UTF-8?q?=E6=9F=90=E4=B8=AA=E5=85=AB=E5=98=8E?= =?UTF-8?q?=E7=BD=91=E7=AB=99=E4=B8=8D=E4=BC=9A=E8=87=AA=E5=B7=B1=E4=B8=8B?= =?UTF-8?q?=EF=BC=88?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index b1ad02d..68f5e49 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,4 @@ fastapi~=0.116.1 jmcomic~=2.6.10 uvicorn~=0.35.0 gunicorn~=23.0.0 +httpx From cd46ce40f506ccb7ed6a404841276b668df5edc0 Mon Sep 17 00:00:00 2001 From: Maverick Charmer Date: Fri, 31 Oct 2025 04:01:26 +0800 Subject: [PATCH 3/6] =?UTF-8?q?=E4=BD=BFtest=5Fapi=E5=8F=AF=E4=BB=A5?= =?UTF-8?q?=E4=BB=8Emain=E6=9E=84=E5=BB=BAapp?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/test_api.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test_api.py b/tests/test_api.py index c07f8b9..1d9ecf6 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -2,6 +2,11 @@ import time import os from pathlib import Path +import sys + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + from main import app current_dir = os.getcwd() From 0162feee0c595bf2d40b80407d438531de684063 Mon Sep 17 00:00:00 2001 From: Maverick Charmer Date: Fri, 31 Oct 2025 04:09:26 +0800 Subject: [PATCH 4/6] Update --- tests/test_api.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tests/test_api.py b/tests/test_api.py index 1d9ecf6..c07f8b9 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -2,11 +2,6 @@ import time import os from pathlib import Path -import sys - -ROOT = Path(__file__).resolve().parents[1] -sys.path.insert(0, str(ROOT)) - from main import app current_dir = os.getcwd() From 402879d62cbbe65d8f35370375d22c0e7c7ba49a Mon Sep 17 00:00:00 2001 From: Maverick Charmer Date: Fri, 31 Oct 2025 04:21:24 +0800 Subject: [PATCH 5/6] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E6=B5=8B=E8=AF=95?= =?UTF-8?q?=E5=AE=8C=E6=88=90=E5=90=8E=E5=88=A0=E9=99=A4temp=E6=96=87?= =?UTF-8?q?=E4=BB=B6=E5=A4=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pytest.ini | 3 +++ tests/test_api.py | 3 +++ 2 files changed, 6 insertions(+) create mode 100644 pytest.ini diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..decb2e8 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +pythonpath = . +testpaths = tests \ No newline at end of file diff --git a/tests/test_api.py b/tests/test_api.py index c07f8b9..457774a 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,3 +1,4 @@ +import shutil from fastapi.testclient import TestClient import time import os @@ -78,3 +79,5 @@ def test_download_album(): assert zip_file_path.exists() == True zip_file_path.unlink() assert zip_file_path.exists() == False + +shutil.rmtree(FILE_PATH) \ No newline at end of file From 576280aa05a8aefb5974df5ccc91e71439f62c19 Mon Sep 17 00:00:00 2001 From: Maverick Charmer Date: Fri, 31 Oct 2025 04:33:41 +0800 Subject: [PATCH 6/6] =?UTF-8?q?=E6=B7=BB=E5=8A=A0pytest-timeout?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pytest.ini | 3 ++- requirements.txt | 5 +++-- tests/test_api.py | 6 ++---- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/pytest.ini b/pytest.ini index decb2e8..f58d47b 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,3 +1,4 @@ [pytest] pythonpath = . -testpaths = tests \ No newline at end of file +testpaths = tests +addopts = -vv -ra -s --maxfail=1 --durations=10 --timeout=60 --timeout-method=thread \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 68f5e49..c2ee644 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ starlette -fastapi~=0.116.1 +fastapi jmcomic~=2.6.10 -uvicorn~=0.35.0 +uvicorn~=0.38.0 gunicorn~=23.0.0 httpx +pytest-timeout \ No newline at end of file diff --git a/tests/test_api.py b/tests/test_api.py index 457774a..58eb3e2 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,4 +1,3 @@ -import shutil from fastapi.testclient import TestClient import time import os @@ -9,6 +8,7 @@ FILE_PATH = Path(f"{current_dir}/temp") os.makedirs(FILE_PATH, exist_ok=True) + def test_read_root(): client = TestClient(app) nowtimestamp = int(time.time() * 1000) @@ -78,6 +78,4 @@ def test_download_album(): zip_file_path = FILE_PATH / zip_file_name assert zip_file_path.exists() == True zip_file_path.unlink() - assert zip_file_path.exists() == False - -shutil.rmtree(FILE_PATH) \ No newline at end of file + assert zip_file_path.exists() == False \ No newline at end of file