Skip to content
Closed

Dev #16

Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
114 changes: 44 additions & 70 deletions main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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
Comment on lines +59 to +61
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Issue: Inadequate Error Handling in WebSocket Closure

The exception handling in the _send_and_close method is too generic, catching all exceptions and passing silently if any occur during the WebSocket closure. This approach can obscure underlying issues that might need attention.

Recommendation:

  • Modify the exception handling to log the specific error or rethrow it after logging. This will help in diagnosing issues during WebSocket operations.
try:
    await websocket.close()
except Exception as e:
    print(f"[WebSocket] Error closing connection for client {client_id}: {e}")
    raise e

self.active_connections.pop(client_id, None)
else:
print(f"[WebSocket Error] 客户端 {client_id} 的连接不存在。")

print(f"[WebSocket] 未找到客户端 {client_id} 的连接,无法发送消息。")

manager = ConnectionManager()

Expand Down Expand Up @@ -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)


# --- 阻塞任务处理函数 (在新线程中运行) ---
Expand All @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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 文件下载路由 ---
Comment on lines 153 to 207
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Issue: Lack of Modularity in Task Execution

The function sync_download_and_zip_task handles multiple responsibilities: downloading, zipping, and sending notifications. This lack of modularity can make the code harder to maintain and test.

Recommendation:

  • Refactor the function into smaller, more focused functions. For example, separate the downloading and zipping into one function and the notification sending into another. This separation can improve the clarity and maintainability of the code.
async def perform_download_and_zip(album_id, option):
    # download and zip logic here

async def send_notification(client_id, message):
    # notification sending logic here

Expand All @@ -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(
Expand Down
4 changes: 4 additions & 0 deletions pytest.ini
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
[pytest]
pythonpath = .
testpaths = tests
addopts = -vv -ra -s --maxfail=1 --durations=10 --timeout=60 --timeout-method=thread
6 changes: 4 additions & 2 deletions requirements.txt
Original file line number Diff line number Diff line change
@@ -1,5 +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
81 changes: 81 additions & 0 deletions tests/test_api.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
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)
Comment on lines +8 to +9
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The directory creation at line 10 does not handle potential exceptions that could arise if the directory cannot be created. This could lead to unhandled exceptions later in the tests when trying to write to this directory.

Recommended Solution:
Wrap the os.makedirs call in a try-except block to handle potential OSError exceptions and add appropriate error handling or logging to manage these cases effectively.



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:
Comment on lines +40 to +62
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hardcoded values such as aid and client_id are used in the test cases. This reduces the flexibility and robustness of the tests, making them potentially fail if the underlying data changes or in different environments where these IDs might not be valid.

Recommended Solution:
Replace hardcoded values with dynamically generated or configured values. This can be achieved by fetching these values from a configuration file, environment variables, or setup methods, ensuring that the tests remain valid and reliable under different conditions.

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"文件 '[酸菜鱼ゅ°]ヒルチャールに败北した胡桃 表情、台词差分' 已完成处理,可以下载。"
}
Comment on lines +62 to +74
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The WebSocket connection and data reception in this test case do not include error handling for cases where the WebSocket connection might fail or the data format received is not as expected. This can lead to the test failing without clear diagnostics or in an uncontrolled manner.

Recommended Solution:
Implement error handling around the WebSocket operations. Use try-except blocks to catch exceptions related to WebSocket connections and validate the format of received data to ensure it meets expected criteria before proceeding with assertions.

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
Comment on lines +12 to +81
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The TestClient is instantiated in every test function which is inefficient and can be optimized. Repeated instantiation can be avoided by using a fixture to create the client once and reuse it across all tests.

Recommended Solution:
Use a pytest fixture to instantiate the TestClient and pass it to each test function. This will not only improve the performance by reducing the setup time for each test but also make the code cleaner and easier to maintain.

5 changes: 4 additions & 1 deletion tests/test_jmcomic_lib.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,15 +6,17 @@
FILE_PATH = Path(f"{current_dir}/temp")
os.makedirs(FILE_PATH, exist_ok=True)
Comment on lines 6 to 7
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hardcoded File Path Issue

The FILE_PATH is constructed using a hardcoded relative path which might lead to issues in different environments where the directory structure or permissions differ. This can affect the portability and robustness of the tests.

Recommendation:
Consider using a temporary directory that is managed by Python's tempfile module, which ensures that the directory is correctly handled across different operating systems and environments. Replace line 6 with:

import tempfile
FILE_PATH = Path(tempfile.mkdtemp())



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)
Comment on lines 10 to 22
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Test Dependency on External Services

The functions test_get_comic_info and test_rank_comic depend on the availability and responsiveness of external services to fetch comic data. This can lead to flaky tests if the external service is down or the data changes unexpectedly.

Recommendation:
To improve the reliability of tests, consider mocking the responses from jmcomic.JmOption.default().new_jm_client() using a library like unittest.mock. This approach allows you to control the data returned by the client and ensures that your tests are not affected by external changes or downtimes.

Expand All @@ -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:
Expand Down