ajyy/backend/tests/test_upload_validate.py

import pytest
from app.core.response import AppHTTPError
from app.core.upload_validate import validate_upload, validate_image_upload

PNG = b"\x89PNG\r\n\x1a\n" + b"\x00" * 16
JPEG = b"\xff\xd8\xff" + b"\x00" * 16
PDF = b"%PDF-1.4" + b"\x00" * 16
EXE_MZ = b"MZ" + b"\x00" * 64


def test_rejects_exe_extension():
    with pytest.raises(AppHTTPError) as ei:
        validate_upload("malware.exe", EXE_MZ)
    assert ei.value.status_code == 400


def test_rejects_png_extension_with_exe_magic():
    with pytest.raises(AppHTTPError) as ei:
        validate_upload("fake.png", EXE_MZ)
    assert ei.value.status_code == 400
    assert "类型" in ei.value.message or "格式" in ei.value.message or "内容" in ei.value.message


def test_accepts_real_png():
    validate_upload("a.png", PNG)


def test_accepts_jpeg_and_pdf():
    validate_upload("a.jpg", JPEG)
    validate_upload("a.pdf", PDF)


def test_image_upload_rejects_pdf():
    with pytest.raises(AppHTTPError):
        validate_image_upload("x.pdf", PDF)


def test_image_upload_accepts_png():
    validate_image_upload("x.png", PNG)