ai_pipeline/tests/pipeline/service/test_file.py

59 lines
1.6 KiB
Python

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()