添加了对API的测试;修改了WS的管理类

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
 

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(

requirements.txt

@@ -3,3 +3,4 @@ fastapi~=0.116.1
 jmcomic~=2.6.10
 uvicorn~=0.35.0
 gunicorn~=23.0.0
+httpx

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

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: