feat: test

This commit is contained in:
李如威 2026-01-09 16:33:34 +08:00
parent 775469a37f
commit bbf8ccfb4a
8 changed files with 251 additions and 31 deletions

View File

@ -1,2 +1,5 @@
[pytest]
pythonpath = .
pythonpath = .
asyncio_mode = auto
asyncio_default_fixture_loop_scope = function
filterwarnings = ignore:.*Swig.*:DeprecationWarning

View File

@ -1,4 +1,4 @@
#!/usr/bin/env bash
export $(cat .env | xargs)
pytest -s -W ignore::DeprecationWarning tests/pipeline/api/test_rag.py::test_chat_message_list
pytest -s -W ignore::DeprecationWarning tests/pipeline/api/test_rag.py::test_chat_message_list_async

View File

@ -1,8 +1,11 @@
from pydantic_settings import BaseSettings
from pydantic import ConfigDict
from pydantic import MySQLDsn
class Settings(BaseSettings):
model_config = ConfigDict(env_file=".env")
logger_level: str
env: str
version: str
@ -23,8 +26,7 @@ class Settings(BaseSettings):
es_user: str
es_password: str
mysql_dsn: MySQLDsn
class Config:
env_file = ".env"
# def _read_config() -> Settings:

View File

@ -19,7 +19,12 @@ async def lifespan(app: FastAPI):
await close_tortoise()
app = FastAPI(title="AI Pipeline", description="轻量级 AI Pipeline", version=config.version, lifespan=lifespan)
app = FastAPI(
title="AI Pipeline",
description="轻量级 AI Pipeline",
version=config.version,
lifespan=lifespan,
)
include_router(app)

View File

@ -5,9 +5,10 @@ from uuid import uuid4
class BaseModel(Model):
id = fields.IntField(pk=True)
id = fields.IntField(primary_key=True)
create_at = fields.DatetimeField(auto_now_add=True)
update_at = fields.DatetimeField(auto_now=True)
deleted_at = fields.DatetimeField(null=True)
is_deleted = fields.IntField(default=0)
class Meta:
@ -74,18 +75,18 @@ class User(BaseModel):
class Chat(BaseModel):
user_id = fields.IntField(null=True)
user_id = fields.IntField(default=0, db_index=True)
title = fields.CharField(max_length=50)
uuid = fields.UUIDField(index=True)
uuid = fields.UUIDField(db_index=True)
class Meta:
table = "tb_chats"
class Message(BaseModel):
user_id = fields.IntField(null=True)
chat_id = fields.IntField(null=True, index=True)
uuid = fields.UUIDField(index=True)
user_id = fields.IntField(default=0, db_index=True)
chat_id = fields.IntField(null=True, db_index=True)
uuid = fields.UUIDField(db_index=True)
content = fields.JSONField(null=True)
class Meta:
@ -98,3 +99,22 @@ class Message(BaseModel):
if "answer" in self.content and self.content["answer"]:
llm_messages.append({"role": "assistant", "content": self.content["answer"]})
return llm_messages
class File(BaseModel):
user_id = fields.IntField(default=0, db_index=True)
uuid = fields.UUIDField(db_index=True)
# 文件基础属性
origin_name = fields.CharField(max_length=244, description="原始名字")
stored_name = fields.CharField(max_length=244, description="存储名字")
stored_path = fields.CharField(max_length=500, description="存储相对路径(不包含文件名字)")
ext = fields.CharField(max_length=20, description="文件扩展名")
mine_type = fields.CharField(max_length=100, description="MIME类型")
size = fields.IntField(default=0, description="文件大小(字节)")
hash_md5 = fields.CharField(max_length=32, description="MD5", db_index=True)
# 业务属性
status = fields.IntField(default=1, description="状态:1正常 2禁用 3删除")
biz_type = fields.CharField(max_length=50, description="业务类型", null=True)
class Meta:
table = "tb_file"

View File

@ -0,0 +1,150 @@
import os
import hashlib
from uuid import uuid4
from datetime import datetime, timezone
from typing import Optional
import aiofiles
import aiofiles.os as aioos
from src.pipeline.models import File as TFile
class FileService:
def __init__(self, base_path: str = "files/uploads"):
"""
base_path: 存储根目录默认当前工作目录下的 storage/uploads
"""
self.base_path = base_path
os.makedirs(self.base_path, exist_ok=True)
def _chunk_dir(self, upload_id: str) -> str:
return os.path.join(self.base_path, "chunks", upload_id)
def _final_dir(self, when: datetime) -> str:
return os.path.join(self.base_path, when.strftime("%Y"), when.strftime("%m"), when.strftime("%d"))
async def save_chunk(self, upload_id: str, chunk_index: int, chunk_bytes: bytes) -> str:
"""
保存一个分片到临时目录异步
返回保存的分片路径绝对
"""
chunk_dir = self._chunk_dir(upload_id)
# 使用同步创建目录是可以接受的,但也支持异步创建
await aioos.makedirs(chunk_dir, exist_ok=True)
part_path = os.path.join(chunk_dir, f"part_{chunk_index}")
async with aiofiles.open(part_path, "wb") as f:
await f.write(chunk_bytes)
return part_path
async def _async_rmtree(self, path: str) -> None:
"""异步递归删除目录及其内容。"""
try:
entries = await aioos.listdir(path)
except FileNotFoundError:
return
for entry in entries:
full = os.path.join(path, entry)
try:
await aioos.remove(full)
except IsADirectoryError:
await self._async_rmtree(full)
try:
await aioos.rmdir(full)
except Exception:
pass
except Exception:
# 可能是目录或其他问题,尝试递归删除
await self._async_rmtree(full)
try:
await aioos.rmdir(full)
except Exception:
pass
try:
await aioos.rmdir(path)
except Exception:
pass
async def finalize_upload(
self,
user_id: int,
upload_id: str,
total_chunks: int,
origin_name: str,
mime_type: Optional[str] = None,
biz_type: Optional[str] = None,
) -> TFile:
"""
upload_id 对应的所有分片合并为最终文件异步写入存储并在数据库中创建 File 记录
返回创建的 File ORM 对象
"""
chunk_dir = self._chunk_dir(upload_id)
# 检查分片目录是否存在
try:
await aioos.listdir(chunk_dir)
except FileNotFoundError:
raise FileNotFoundError("chunk directory not found")
# 确定扩展名和目标路径/文件名
_, ext = os.path.splitext(origin_name)
ext = ext.lstrip(".") if ext else ""
now = datetime.now(timezone.utc)
final_dir = self._final_dir(now)
await aioos.makedirs(final_dir, exist_ok=True)
stored_name = f"{uuid4().hex}{('.' + ext) if ext else ''}"
final_path = os.path.join(final_dir, stored_name)
# 合并并计算 md5 与大小(异步读写)
md5 = hashlib.md5()
total_size = 0
async with aiofiles.open(final_path, "wb") as fw:
for i in range(total_chunks):
part_path = os.path.join(chunk_dir, f"part_{i}")
try:
async with aiofiles.open(part_path, "rb") as fr:
while True:
buf = await fr.read(1024 * 1024)
if not buf:
break
await fw.write(buf)
md5.update(buf)
total_size += len(buf)
except FileNotFoundError:
# 清理已写入的文件
try:
await aioos.remove(final_path)
except Exception:
pass
raise FileNotFoundError(f"missing chunk: {i}")
md5_hex = md5.hexdigest()
# 清理分片目录(异步)
await self._async_rmtree(chunk_dir)
# stored_path 应为相对路径(不包含文件名),相对于 base_path
stored_path_rel = os.path.relpath(final_dir, self.base_path)
# 在数据库中创建 File 记录并返回
file_uuid = uuid4()
file_record = await TFile.create(
user_id=user_id,
uuid=file_uuid,
origin_name=origin_name,
stored_name=stored_name,
stored_path=stored_path_rel,
ext=ext,
mine_type=mime_type or "",
size=total_size,
hash_md5=md5_hex,
biz_type=biz_type or None,
)
return file_record
async def abort_upload(self, upload_id: str) -> None:
"""删除临时分片目录(中止上传用,异步)。"""
chunk_dir = self._chunk_dir(upload_id)
await self._async_rmtree(chunk_dir)
file_service = FileService()

View File

@ -2,7 +2,6 @@ from asgi_lifespan import LifespanManager
from httpx import AsyncClient, ASGITransport
from src.pipeline.main import app
from src.pipeline.utils import logger
from fastapi.testclient import TestClient
import pytest
import pytest_asyncio
import json
@ -12,7 +11,7 @@ import json
# ----------------------
@pytest_asyncio.fixture(scope="session")
@pytest_asyncio.fixture
async def app_client_async():
async with LifespanManager(app):
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
@ -21,7 +20,7 @@ async def app_client_async():
@pytest.mark.asyncio
async def test_stream_chat(app_client_async):
res = await app_client.post("/api/v1/rag/stream-chat", json={"query": "hello"})
res = await app_client_async.post("/api/v1/rag/stream-chat", json={"query": "hello"})
assert res.status_code == 200
sse_messages = []
async for line in res.aiter_lines():
@ -43,20 +42,3 @@ async def test_chat_message_list_async(app_client_async):
res = await app_client_async.get("/api/v1/rag/message-list/a1c9108c-9201-4e60-a436-505edec3f47e")
assert res.status_code == 200
logger.debug(res.json())
# ----------------------
# normal
# ----------------------
@pytest.fixture(scope="session")
def app_client():
client = TestClient(app=app)
return client
def test_chat_message_list(app_client):
res = app_client.get("/api/v1/rag/message-list/a1c9108c-9201-4e60-a436-505edec3f47e")
assert res.status_code == 200
logger.debug(res.json())

View File

@ -0,0 +1,58 @@
import hashlib
import os
import pytest
import pytest_asyncio
from src.pipeline.db import Path, init_tortoise, close_tortoise
from src.pipeline.services.file import file_service
from src.pipeline.config import config
from src.pipeline.models import File as TFile
from uuid import uuid4
@pytest_asyncio.fixture
async def init_db():
db_url = str(config.mysql_dsn)
await init_tortoise(db_url=db_url)
yield
await close_tortoise()
@pytest.mark.asyncio
async def test_file_upload(init_db):
upload_id = str(uuid4())
chunks = [b"part1-", b"part2-", b"part3"]
# save chunks
for idx, chunk in enumerate(chunks):
part_path = await file_service.save_chunk(upload_id, idx, chunk)
assert os.path.exists(part_path)
# finalize
origin_name = "hello.txt"
file_record = await file_service.finalize_upload(
user_id=42,
upload_id=upload_id,
total_chunks=len(chunks),
origin_name=origin_name,
mime_type="text/plain",
biz_type="unittest",
)
# DB record
assert isinstance(file_record, TFile)
assert file_record.user_id == 42
assert file_record.origin_name == origin_name
# file exists on disk
stored_path = Path(file_service.base_path) / Path(file_record.stored_path)
final_path = stored_path / file_record.stored_name
assert final_path.exists()
# md5 and size correct
combined = b"".join(chunks)
assert file_record.hash_md5 == hashlib.md5(combined).hexdigest()
assert file_record.size == len(combined)
# chunks directory removed
chunk_dir = Path(file_service.base_path) / "chunks" / upload_id
assert not chunk_dir.exists()