feat: test
This commit is contained in:
parent
775469a37f
commit
bbf8ccfb4a
|
|
@ -1,2 +1,5 @@
|
|||
[pytest]
|
||||
pythonpath = .
|
||||
pythonpath = .
|
||||
asyncio_mode = auto
|
||||
asyncio_default_fixture_loop_scope = function
|
||||
filterwarnings = ignore:.*Swig.*:DeprecationWarning
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
Loading…
Reference in New Issue