59 lines
1.6 KiB
Python
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()
|