feat(knowledge): 知识源维护与 TXT/Markdown/PDF 导入
支持文本、Markdown 与非扫描 PDF 上传解析,展示提取文本与切块结果。
This commit is contained in:
1 parent
4388e1d203
commit
53825b0cfb
23 files changed
+7545
No files matched your search
@@ -0,0 +1,512 @@
|
||||
"""
|
||||
Copyright 2024-2026 ChatterMate
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import base64
|
||||
import os
|
||||
import tempfile
|
||||
from unittest.mock import Mock, patch, AsyncMock, MagicMock
|
||||
from uuid import uuid4
|
||||
|
||||
from app.services.file_upload_service import (
|
||||
FileUploadService,
|
||||
ALLOWED_IMAGE_TYPES,
|
||||
ALLOWED_DOCUMENT_TYPES,
|
||||
ALLOWED_FILE_TYPES,
|
||||
MAX_FILE_SIZE,
|
||||
MAX_IMAGE_SIZE
|
||||
)
|
||||
|
||||
|
||||
class TestFileUploadServiceValidation:
|
||||
"""Test file validation methods"""
|
||||
|
||||
def test_validate_file_type_allowed_image(self):
|
||||
"""Test validation of allowed image types"""
|
||||
for content_type in ALLOWED_IMAGE_TYPES:
|
||||
assert FileUploadService.validate_file_type(content_type) is True
|
||||
|
||||
def test_validate_file_type_allowed_document(self):
|
||||
"""Test validation of allowed document types"""
|
||||
for content_type in ALLOWED_DOCUMENT_TYPES:
|
||||
assert FileUploadService.validate_file_type(content_type) is True
|
||||
|
||||
def test_validate_file_type_disallowed(self):
|
||||
"""Test validation of disallowed file types"""
|
||||
disallowed_types = [
|
||||
'video/mp4',
|
||||
'audio/mp3',
|
||||
'application/x-executable',
|
||||
'text/html',
|
||||
'application/javascript'
|
||||
]
|
||||
for content_type in disallowed_types:
|
||||
assert FileUploadService.validate_file_type(content_type) is False
|
||||
|
||||
def test_validate_file_size_image_within_limit(self):
|
||||
"""Test image file size validation within limits"""
|
||||
# Test image under image size limit
|
||||
image_size = MAX_IMAGE_SIZE - 1000 # 1KB under limit
|
||||
is_valid, error = FileUploadService.validate_file_size(image_size, 'image/jpeg')
|
||||
assert is_valid is True
|
||||
assert error is None
|
||||
|
||||
def test_validate_file_size_image_exceeds_image_limit(self):
|
||||
"""Test image file size validation exceeding image limit"""
|
||||
# Test image over image size limit
|
||||
image_size = MAX_IMAGE_SIZE + 1000 # 1KB over image limit
|
||||
is_valid, error = FileUploadService.validate_file_size(image_size, 'image/jpeg')
|
||||
assert is_valid is False
|
||||
assert "Image file size exceeds maximum" in error
|
||||
assert "5.0MB" in error
|
||||
|
||||
def test_validate_file_size_document_within_limit(self):
|
||||
"""Test document file size validation within limits"""
|
||||
# Test document under general size limit
|
||||
doc_size = MAX_FILE_SIZE - 1000 # 1KB under limit
|
||||
is_valid, error = FileUploadService.validate_file_size(doc_size, 'application/pdf')
|
||||
assert is_valid is True
|
||||
assert error is None
|
||||
|
||||
def test_validate_file_size_document_exceeds_limit(self):
|
||||
"""Test document file size validation exceeding limit"""
|
||||
# Test document over general size limit
|
||||
doc_size = MAX_FILE_SIZE + 1000 # 1KB over limit
|
||||
is_valid, error = FileUploadService.validate_file_size(doc_size, 'application/pdf')
|
||||
assert is_valid is False
|
||||
assert "File size exceeds maximum" in error
|
||||
assert "10.0MB" in error
|
||||
|
||||
def test_validate_file_size_image_exceeds_general_limit(self):
|
||||
"""Test image file size validation exceeding general limit"""
|
||||
# Test image over general size limit (but under image limit is irrelevant)
|
||||
image_size = MAX_FILE_SIZE + 1000 # 1KB over general limit
|
||||
is_valid, error = FileUploadService.validate_file_size(image_size, 'image/png')
|
||||
assert is_valid is False
|
||||
# The service checks image size limit first, so we get image-specific error
|
||||
assert "Image file size exceeds maximum" in error
|
||||
assert "5.0MB" in error
|
||||
|
||||
|
||||
class TestFileUploadServiceUpload:
|
||||
"""Test file upload functionality"""
|
||||
|
||||
@pytest.fixture
|
||||
def sample_file_data(self):
|
||||
"""Create sample file data for testing"""
|
||||
content = b"This is a test file content"
|
||||
encoded_content = base64.b64encode(content).decode('utf-8')
|
||||
return {
|
||||
'content': encoded_content,
|
||||
'filename': 'test_file.txt',
|
||||
'content_type': 'text/plain',
|
||||
'size': len(content)
|
||||
}
|
||||
|
||||
@pytest.fixture
|
||||
def sample_image_data(self):
|
||||
"""Create sample image data for testing"""
|
||||
# Simple fake image content
|
||||
content = b"fake_image_data_" * 100 # Small image
|
||||
encoded_content = base64.b64encode(content).decode('utf-8')
|
||||
return {
|
||||
'content': encoded_content,
|
||||
'filename': 'test_image.jpg',
|
||||
'content_type': 'image/jpeg',
|
||||
'size': len(content)
|
||||
}
|
||||
|
||||
@pytest.fixture
|
||||
def large_file_data(self):
|
||||
"""Create large file data that exceeds size limits"""
|
||||
content = b"x" * (MAX_FILE_SIZE + 1000) # Exceeds max file size
|
||||
encoded_content = base64.b64encode(content).decode('utf-8')
|
||||
return {
|
||||
'content': encoded_content,
|
||||
'filename': 'large_file.pdf',
|
||||
'content_type': 'application/pdf',
|
||||
'size': len(content)
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_s3_success(self, sample_file_data):
|
||||
"""Test successful file upload to S3"""
|
||||
org_id = str(uuid4())
|
||||
user_id = str(uuid4())
|
||||
|
||||
with patch('app.services.file_upload_service.settings') as mock_settings, \
|
||||
patch('app.services.file_upload_service.upload_file_to_s3') as mock_s3_upload, \
|
||||
patch('app.services.file_upload_service.resolve_public_url') as mock_signed_url, \
|
||||
patch('uuid.uuid4') as mock_uuid:
|
||||
|
||||
# Configure mocks
|
||||
mock_settings.S3_FILE_STORAGE = True
|
||||
mock_uuid.return_value = Mock()
|
||||
mock_uuid.return_value.__str__ = Mock(return_value='test-uuid-123')
|
||||
|
||||
mock_s3_upload.return_value = 'https://s3.amazonaws.com/bucket/chat_attachments/org123/test-uuid-123.txt'
|
||||
mock_signed_url.return_value = 'https://s3.amazonaws.com/bucket/signed-url'
|
||||
|
||||
result = await FileUploadService.upload_file(
|
||||
file_data=sample_file_data,
|
||||
org_id=org_id,
|
||||
user_id=user_id
|
||||
)
|
||||
|
||||
# Verify S3 upload was called
|
||||
mock_s3_upload.assert_called_once()
|
||||
call_args = mock_s3_upload.call_args
|
||||
assert call_args[1]['folder'] == f'chat_attachments/{org_id}'
|
||||
assert call_args[1]['filename'] == 'test-uuid-123.txt'
|
||||
assert call_args[1]['content_type'] == 'text/plain'
|
||||
|
||||
# Verify signed URL was generated
|
||||
mock_signed_url.assert_called_once()
|
||||
|
||||
# Verify result
|
||||
assert result['file_url'] == 'https://s3.amazonaws.com/bucket/chat_attachments/org123/test-uuid-123.txt'
|
||||
assert result['signed_url'] == 'https://s3.amazonaws.com/bucket/signed-url'
|
||||
assert result['filename'] == 'test_file.txt'
|
||||
assert result['content_type'] == 'text/plain'
|
||||
assert result['size'] == len(base64.b64decode(sample_file_data['content']))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_local_success(self, sample_file_data):
|
||||
"""Test successful file upload to local storage"""
|
||||
org_id = str(uuid4())
|
||||
customer_id = str(uuid4())
|
||||
|
||||
with patch('app.services.file_upload_service.settings') as mock_settings, \
|
||||
patch('uuid.uuid4') as mock_uuid, \
|
||||
patch('os.makedirs') as mock_makedirs, \
|
||||
patch('builtins.open', create=True) as mock_open:
|
||||
|
||||
# Configure mocks
|
||||
mock_settings.S3_FILE_STORAGE = False
|
||||
mock_uuid.return_value = Mock()
|
||||
mock_uuid.return_value.__str__ = Mock(return_value='test-uuid-456')
|
||||
|
||||
# Mock file operations
|
||||
mock_file = MagicMock()
|
||||
mock_open.return_value.__enter__.return_value = mock_file
|
||||
|
||||
result = await FileUploadService.upload_file(
|
||||
file_data=sample_file_data,
|
||||
org_id=org_id,
|
||||
customer_id=customer_id
|
||||
)
|
||||
|
||||
# Verify directory creation
|
||||
mock_makedirs.assert_called_once_with(f'uploads/chat_attachments/{org_id}', exist_ok=True)
|
||||
|
||||
# Verify file write
|
||||
mock_open.assert_called_once_with(f'uploads/chat_attachments/{org_id}/test-uuid-456.txt', 'wb')
|
||||
mock_file.write.assert_called_once()
|
||||
|
||||
# Verify result
|
||||
assert result['file_url'] == f'/uploads/chat_attachments/{org_id}/test-uuid-456.txt'
|
||||
# Local attachments are never static-served: clients get a signed link.
|
||||
assert result['signed_url'].startswith(
|
||||
f'/api/v1/files/local/chat_attachments/{org_id}/test-uuid-456.txt?exp=')
|
||||
assert result['filename'] == 'test_file.txt'
|
||||
assert result['content_type'] == 'text/plain'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_missing_content(self):
|
||||
"""Test upload with missing content"""
|
||||
file_data = {
|
||||
'filename': 'test.txt',
|
||||
'content_type': 'text/plain'
|
||||
# Missing 'content'
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
await FileUploadService.upload_file(
|
||||
file_data=file_data,
|
||||
org_id=str(uuid4())
|
||||
)
|
||||
|
||||
assert "Missing required file data" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_missing_filename(self):
|
||||
"""Test upload with missing filename"""
|
||||
file_data = {
|
||||
'content': base64.b64encode(b"test content").decode('utf-8'),
|
||||
'content_type': 'text/plain'
|
||||
# Missing 'filename'
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
await FileUploadService.upload_file(
|
||||
file_data=file_data,
|
||||
org_id=str(uuid4())
|
||||
)
|
||||
|
||||
assert "Missing required file data" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_invalid_type(self, sample_file_data):
|
||||
"""Test upload with invalid file type"""
|
||||
# Modify to invalid type
|
||||
sample_file_data['content_type'] = 'video/mp4'
|
||||
sample_file_data['filename'] = 'test.mp4'
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
await FileUploadService.upload_file(
|
||||
file_data=sample_file_data,
|
||||
org_id=str(uuid4())
|
||||
)
|
||||
|
||||
# Error should indicate the file type is not allowed
|
||||
assert "not allowed" in str(exc_info.value) or "File type" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_invalid_base64(self):
|
||||
"""Test upload with invalid base64 content"""
|
||||
file_data = {
|
||||
'content': 'Z', # Invalid base64 - single character that can't be decoded properly
|
||||
'filename': 'test.txt',
|
||||
'content_type': 'text/plain'
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
await FileUploadService.upload_file(
|
||||
file_data=file_data,
|
||||
org_id=str(uuid4())
|
||||
)
|
||||
|
||||
assert "Invalid file content encoding" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_size_exceeds_limit(self, large_file_data):
|
||||
"""Test upload with file size exceeding limits"""
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
await FileUploadService.upload_file(
|
||||
file_data=large_file_data,
|
||||
org_id=str(uuid4())
|
||||
)
|
||||
|
||||
assert "File size exceeds maximum" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_image_size_exceeds_limit(self):
|
||||
"""Test upload with image size exceeding image limits"""
|
||||
# Create large image that exceeds image size limit
|
||||
content = b"x" * (MAX_IMAGE_SIZE + 1000)
|
||||
encoded_content = base64.b64encode(content).decode('utf-8')
|
||||
|
||||
file_data = {
|
||||
'content': encoded_content,
|
||||
'filename': 'large_image.jpg',
|
||||
'content_type': 'image/jpeg'
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
await FileUploadService.upload_file(
|
||||
file_data=file_data,
|
||||
org_id=str(uuid4())
|
||||
)
|
||||
|
||||
assert "Image file size exceeds maximum" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_s3_upload_failure(self, sample_file_data):
|
||||
"""Test handling of S3 upload failure"""
|
||||
with patch('app.services.file_upload_service.settings') as mock_settings, \
|
||||
patch('app.services.file_upload_service.upload_file_to_s3') as mock_s3_upload:
|
||||
|
||||
mock_settings.S3_FILE_STORAGE = True
|
||||
mock_s3_upload.side_effect = Exception("S3 upload failed")
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
await FileUploadService.upload_file(
|
||||
file_data=sample_file_data,
|
||||
org_id=str(uuid4())
|
||||
)
|
||||
|
||||
assert "Failed to upload file" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_local_write_failure(self, sample_file_data):
|
||||
"""Test handling of local file write failure"""
|
||||
with patch('app.services.file_upload_service.settings') as mock_settings, \
|
||||
patch('os.makedirs') as mock_makedirs, \
|
||||
patch('builtins.open', create=True) as mock_open:
|
||||
|
||||
mock_settings.S3_FILE_STORAGE = False
|
||||
mock_open.side_effect = OSError("Permission denied")
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
await FileUploadService.upload_file(
|
||||
file_data=sample_file_data,
|
||||
org_id=str(uuid4())
|
||||
)
|
||||
|
||||
assert "Failed to upload file" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_default_content_type(self):
|
||||
"""Test upload with default content type when not provided"""
|
||||
content = b"test content"
|
||||
file_data = {
|
||||
'content': base64.b64encode(content).decode('utf-8'),
|
||||
'filename': 'test.txt'
|
||||
# No content_type provided - will default to 'application/octet-stream' which is not allowed
|
||||
}
|
||||
|
||||
# Should raise error because default content type is not allowed
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
await FileUploadService.upload_file(
|
||||
file_data=file_data,
|
||||
org_id=str(uuid4())
|
||||
)
|
||||
|
||||
# Error could be extension mismatch or file type not allowed
|
||||
error_str = str(exc_info.value).lower()
|
||||
assert "not allowed" in error_str or "does not match" in error_str or "file type" in error_str
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_folder_structure_user(self, sample_file_data):
|
||||
"""Test folder structure for user uploads"""
|
||||
org_id = str(uuid4())
|
||||
user_id = str(uuid4())
|
||||
|
||||
with patch('app.services.file_upload_service.settings') as mock_settings, \
|
||||
patch('app.services.file_upload_service.upload_file_to_s3') as mock_s3_upload, \
|
||||
patch('app.services.file_upload_service.resolve_public_url') as mock_signed_url:
|
||||
|
||||
mock_settings.S3_FILE_STORAGE = True
|
||||
mock_s3_upload.return_value = 'https://s3.amazonaws.com/bucket/file.txt'
|
||||
mock_signed_url.return_value = 'https://s3.amazonaws.com/bucket/signed'
|
||||
|
||||
await FileUploadService.upload_file(
|
||||
file_data=sample_file_data,
|
||||
org_id=org_id,
|
||||
user_id=user_id
|
||||
)
|
||||
|
||||
# Verify folder structure
|
||||
call_args = mock_s3_upload.call_args
|
||||
assert call_args[1]['folder'] == f'chat_attachments/{org_id}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_folder_structure_customer(self, sample_file_data):
|
||||
"""Test folder structure for customer uploads"""
|
||||
org_id = str(uuid4())
|
||||
customer_id = str(uuid4())
|
||||
|
||||
with patch('app.services.file_upload_service.settings') as mock_settings, \
|
||||
patch('app.services.file_upload_service.upload_file_to_s3') as mock_s3_upload, \
|
||||
patch('app.services.file_upload_service.resolve_public_url') as mock_signed_url:
|
||||
|
||||
mock_settings.S3_FILE_STORAGE = True
|
||||
mock_s3_upload.return_value = 'https://s3.amazonaws.com/bucket/file.txt'
|
||||
mock_signed_url.return_value = 'https://s3.amazonaws.com/bucket/signed'
|
||||
|
||||
await FileUploadService.upload_file(
|
||||
file_data=sample_file_data,
|
||||
org_id=org_id,
|
||||
customer_id=customer_id
|
||||
)
|
||||
|
||||
# Verify folder structure (same as user for now)
|
||||
call_args = mock_s3_upload.call_args
|
||||
assert call_args[1]['folder'] == f'chat_attachments/{org_id}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_unique_filename_generation(self, sample_file_data):
|
||||
"""Test that unique filenames are generated"""
|
||||
org_id = str(uuid4())
|
||||
|
||||
with patch('app.services.file_upload_service.settings') as mock_settings, \
|
||||
patch('app.services.file_upload_service.upload_file_to_s3') as mock_s3_upload, \
|
||||
patch('app.services.file_upload_service.resolve_public_url') as mock_signed_url, \
|
||||
patch('uuid.uuid4') as mock_uuid:
|
||||
|
||||
mock_settings.S3_FILE_STORAGE = True
|
||||
mock_s3_upload.return_value = 'https://s3.amazonaws.com/bucket/file.txt'
|
||||
mock_signed_url.return_value = 'https://s3.amazonaws.com/bucket/signed'
|
||||
|
||||
# Mock UUID generation
|
||||
mock_uuid.return_value = Mock()
|
||||
mock_uuid.return_value.__str__ = Mock(return_value='unique-id-789')
|
||||
|
||||
await FileUploadService.upload_file(
|
||||
file_data=sample_file_data,
|
||||
org_id=org_id
|
||||
)
|
||||
|
||||
# Verify unique filename was generated
|
||||
call_args = mock_s3_upload.call_args
|
||||
assert call_args[1]['filename'] == 'unique-id-789.txt'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_no_extension(self):
|
||||
"""Test upload with filename that has no extension"""
|
||||
content = b"test content"
|
||||
file_data = {
|
||||
'content': base64.b64encode(content).decode('utf-8'),
|
||||
'filename': 'test_file_no_extension',
|
||||
'content_type': 'text/plain'
|
||||
}
|
||||
|
||||
# Should raise error because file has no extension
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
await FileUploadService.upload_file(
|
||||
file_data=file_data,
|
||||
org_id=str(uuid4())
|
||||
)
|
||||
|
||||
assert "must have an extension" in str(exc_info.value) or "extension" in str(exc_info.value).lower()
|
||||
|
||||
|
||||
class TestFileUploadServiceConstants:
|
||||
"""Test service constants and configurations"""
|
||||
|
||||
def test_allowed_file_types_combination(self):
|
||||
"""Test that ALLOWED_FILE_TYPES is correct combination"""
|
||||
expected_types = ALLOWED_IMAGE_TYPES | ALLOWED_DOCUMENT_TYPES
|
||||
assert ALLOWED_FILE_TYPES == expected_types
|
||||
|
||||
def test_file_size_constants(self):
|
||||
"""Test file size constants are reasonable"""
|
||||
assert MAX_FILE_SIZE == 10 * 1024 * 1024 # 10MB
|
||||
assert MAX_IMAGE_SIZE == 5 * 1024 * 1024 # 5MB
|
||||
assert MAX_IMAGE_SIZE < MAX_FILE_SIZE # Image limit should be smaller
|
||||
|
||||
def test_allowed_image_types_content(self):
|
||||
"""Test allowed image types include common formats (SVG excluded for XSS safety)"""
|
||||
expected_image_types = {
|
||||
'image/jpeg', 'image/jpg', 'image/png',
|
||||
'image/gif', 'image/webp'
|
||||
}
|
||||
assert ALLOWED_IMAGE_TYPES == expected_image_types
|
||||
# Verify SVG is explicitly NOT included
|
||||
assert 'image/svg+xml' not in ALLOWED_IMAGE_TYPES
|
||||
|
||||
def test_allowed_document_types_content(self):
|
||||
"""Test allowed document types include common formats"""
|
||||
expected_doc_types = {
|
||||
'application/pdf',
|
||||
'application/msword',
|
||||
'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
|
||||
'text/plain',
|
||||
'text/csv',
|
||||
'application/vnd.ms-excel',
|
||||
'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
|
||||
}
|
||||
assert ALLOWED_DOCUMENT_TYPES == expected_doc_types
|
||||
@@ -0,0 +1,94 @@
|
||||
"""
|
||||
Copyright 2024-2026 ChatterMate
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import io
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from PIL import Image
|
||||
|
||||
from app.services.image_security import sanitize_image
|
||||
|
||||
RASTER = {"image/png", "image/jpeg", "image/webp"}
|
||||
|
||||
|
||||
def _img_bytes(fmt: str, size=(64, 48), color=(200, 30, 30)) -> bytes:
|
||||
buf = io.BytesIO()
|
||||
Image.new("RGB", size, color).save(buf, format=fmt)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def test_png_reencoded_to_png():
|
||||
out, ctype, ext = sanitize_image(
|
||||
_img_bytes("PNG"), allowed_content_types=RASTER, max_bytes=2_000_000, max_dim=4000
|
||||
)
|
||||
assert ctype == "image/png" and ext == ".png"
|
||||
assert Image.open(io.BytesIO(out)).format == "PNG"
|
||||
|
||||
|
||||
def test_jpeg_stays_jpeg():
|
||||
out, ctype, ext = sanitize_image(
|
||||
_img_bytes("JPEG"), allowed_content_types=RASTER, max_bytes=2_000_000, max_dim=4000
|
||||
)
|
||||
assert ctype == "image/jpeg" and ext == ".jpg"
|
||||
assert Image.open(io.BytesIO(out)).format == "JPEG"
|
||||
|
||||
|
||||
def test_fit_downscales_longest_side():
|
||||
out, _c, _e = sanitize_image(
|
||||
_img_bytes("PNG", size=(1200, 600)),
|
||||
allowed_content_types=RASTER,
|
||||
max_bytes=2_000_000,
|
||||
max_dim=4000,
|
||||
fit=512,
|
||||
)
|
||||
w, h = Image.open(io.BytesIO(out)).size
|
||||
assert max(w, h) == 512 and (w, h) == (512, 256)
|
||||
|
||||
|
||||
def test_rejects_non_image():
|
||||
with pytest.raises(HTTPException) as e:
|
||||
sanitize_image(b"not an image", allowed_content_types=RASTER, max_bytes=2_000_000, max_dim=4000)
|
||||
assert e.value.status_code == 400
|
||||
|
||||
|
||||
def test_rejects_oversized_bytes():
|
||||
with pytest.raises(HTTPException) as e:
|
||||
sanitize_image(_img_bytes("PNG"), allowed_content_types=RASTER, max_bytes=10, max_dim=4000)
|
||||
assert e.value.status_code == 400
|
||||
|
||||
|
||||
def test_rejects_dimensions_over_cap():
|
||||
with pytest.raises(HTTPException) as e:
|
||||
sanitize_image(
|
||||
_img_bytes("PNG", size=(300, 300)),
|
||||
allowed_content_types=RASTER,
|
||||
max_bytes=2_000_000,
|
||||
max_dim=100,
|
||||
)
|
||||
assert e.value.status_code == 400
|
||||
|
||||
|
||||
def test_rejects_disallowed_format():
|
||||
# GIF is a valid image but not in the allowlist here → rejected.
|
||||
with pytest.raises(HTTPException) as e:
|
||||
sanitize_image(_img_bytes("GIF"), allowed_content_types=RASTER, max_bytes=2_000_000, max_dim=4000)
|
||||
assert e.value.status_code == 400
|
||||
|
||||
|
||||
def test_empty_rejected():
|
||||
with pytest.raises(HTTPException):
|
||||
sanitize_image(b"", allowed_content_types=RASTER, max_bytes=2_000_000, max_dim=4000)
|
||||
@@ -0,0 +1,88 @@
|
||||
"""
|
||||
Copyright 2024-2026 ChatterMate
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import io
|
||||
import os
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from fastapi import UploadFile
|
||||
from PIL import Image
|
||||
|
||||
from app.api.agent import save_file
|
||||
from app.api.users import save_upload_file
|
||||
from app.core.config import settings
|
||||
from app.services.file_storage import local_upload_path, store_upload
|
||||
|
||||
MOUNT = f"{settings.API_V1_STR}/uploads"
|
||||
|
||||
|
||||
def _png_upload(name: str = "a.png") -> UploadFile:
|
||||
buf = io.BytesIO()
|
||||
Image.new("RGB", (2, 2), "red").save(buf, format="PNG")
|
||||
buf.seek(0)
|
||||
return UploadFile(filename=name, file=buf, headers={"content-type": "image/png"})
|
||||
|
||||
|
||||
class TestLocalUploadPath:
|
||||
"""The stored-URL → filesystem-path direction has to read every shape the
|
||||
app has written, because rows predating the normalisation are still bare."""
|
||||
|
||||
def test_current_prefixed_shape(self):
|
||||
assert local_upload_path(f"{MOUNT}/agents/org/a.png") == os.path.join(
|
||||
"uploads", "agents/org/a.png")
|
||||
|
||||
def test_legacy_bare_shape(self):
|
||||
# What agent photos and profile pictures carried before the migration.
|
||||
assert local_upload_path("/uploads/agents/org/a.png") == os.path.join(
|
||||
"uploads", "agents/org/a.png")
|
||||
|
||||
def test_relative_key(self):
|
||||
assert local_upload_path("agents/org/a.png") == os.path.join(
|
||||
"uploads", "agents/org/a.png")
|
||||
|
||||
def test_never_doubles_the_uploads_segment(self):
|
||||
for stored in (f"{MOUNT}/x.png", "/uploads/x.png", "uploads/x.png", "x.png"):
|
||||
assert local_upload_path(stored) == os.path.join("uploads", "x.png"), stored
|
||||
|
||||
|
||||
class TestWritersAgree:
|
||||
"""Every local writer must emit the same prefix, or the frontend resolves
|
||||
the path against the API origin and 404s off the static mount."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_upload_prefixes(self, tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.setattr(settings, "S3_FILE_STORAGE", False)
|
||||
stored = await store_upload(b"x", "help_center", "a.png", "image/png")
|
||||
assert stored.startswith(f"{MOUNT}/")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_photo_prefixes_and_round_trips(self, tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.setattr(settings, "S3_FILE_STORAGE", False)
|
||||
stored = await save_file(_png_upload(), uuid4())
|
||||
assert stored.startswith(f"{MOUNT}/")
|
||||
# The delete path reads it back through local_upload_path.
|
||||
assert os.path.exists(local_upload_path(stored))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_profile_pic_prefixes_and_round_trips(self, tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.setattr(settings, "S3_FILE_STORAGE", False)
|
||||
stored = await save_upload_file(_png_upload(), str(uuid4()), str(uuid4()))
|
||||
assert stored.startswith(f"{MOUNT}/")
|
||||
assert os.path.exists(local_upload_path(stored))
|
||||
Reference in new issue
Block a user