import base64 import hashlib import os from unittest.mock import MagicMock, patch import pytest from sqlalchemy import Engine from sqlalchemy.orm import Session, sessionmaker from werkzeug.exceptions import NotFound from configs import dify_config from models.enums import CreatorUserRole from models.model import Account, EndUser, UploadFile from services.errors.file import BlockedFileExtensionError, FileTooLargeError, UnsupportedFileTypeError from services.file_service import FileService class TestFileService: @pytest.fixture def mock_db_session(self): session = MagicMock(spec=Session) # Mock context manager behavior session.__enter__.return_value = session return session @pytest.fixture def mock_session_maker(self, mock_db_session): maker = MagicMock(spec=sessionmaker) maker.return_value = mock_db_session return maker @pytest.fixture def file_service(self, mock_session_maker): return FileService(session_factory=mock_session_maker) def test_init_with_engine(self): engine = MagicMock(spec=Engine) service = FileService(session_factory=engine) assert isinstance(service._session_maker, sessionmaker) def test_init_with_sessionmaker(self): maker = MagicMock(spec=sessionmaker) service = FileService(session_factory=maker) assert service._session_maker == maker def test_init_invalid_factory(self): with pytest.raises(AssertionError, match="must be a sessionmaker or an Engine."): FileService(session_factory="invalid") @patch("services.file_service.storage") @patch("services.file_service.naive_utc_now") @patch("services.file_service.extract_tenant_id") @patch("services.file_service.file_helpers.get_signed_file_url") def test_upload_file_success( self, mock_get_url, mock_tenant_id, mock_now, mock_storage, file_service: FileService, mock_db_session ): # Setup mock_tenant_id.return_value = "tenant_id" mock_now.return_value = "2024-01-01" mock_get_url.return_value = "http://signed-url" user = MagicMock(spec=Account) user.id = "user_id" content = b"file content" filename = "test.jpg" mimetype = "image/jpeg" # Execute result = file_service.upload_file(filename=filename, content=content, mimetype=mimetype, user=user) # Assert assert isinstance(result, UploadFile) assert result.name == filename assert result.tenant_id == "tenant_id" assert result.size == len(content) assert result.extension == "jpg" assert result.mime_type == mimetype assert result.created_by_role == CreatorUserRole.ACCOUNT assert result.created_by == "user_id" assert result.hash == hashlib.sha3_256(content).hexdigest() assert result.source_url == "http://signed-url" mock_storage.save.assert_called_once() mock_db_session.add.assert_called_once_with(result) mock_db_session.commit.assert_called_once() def test_upload_file_invalid_characters(self, file_service): with pytest.raises(ValueError, match="Filename contains invalid characters"): file_service.upload_file(filename="invalid/file.txt", content=b"", mimetype="text/plain", user=MagicMock()) def test_upload_file_long_filename(self, file_service: FileService, mock_db_session): # Setup long_name = "a" * 210 + ".txt" user = MagicMock(spec=Account) user.id = "user_id" with ( patch("services.file_service.storage"), patch("services.file_service.extract_tenant_id") as mock_tenant, patch("services.file_service.file_helpers.get_signed_file_url"), ): mock_tenant.return_value = "tenant" result = file_service.upload_file(filename=long_name, content=b"test", mimetype="text/plain", user=user) assert len(result.name) <= 205 # 200 + . + extension assert result.name.endswith(".txt") def test_upload_file_blocked_extension(self, file_service): with patch.object(dify_config, "inner_UPLOAD_FILE_EXTENSION_BLACKLIST", "exe"): with pytest.raises(BlockedFileExtensionError): file_service.upload_file( filename="test.exe", content=b"", mimetype="application/octet-stream", user=MagicMock() ) def test_upload_file_unsupported_type_for_datasets(self, file_service): with pytest.raises(UnsupportedFileTypeError): file_service.upload_file( filename="test.jpg", content=b"", mimetype="image/jpeg", user=MagicMock(), source="datasets" ) def test_upload_file_too_large(self, file_service): # 16MB file for an image with 15MB limit content = b"a" * (16 * 1024 * 1024) with patch.object(dify_config, "UPLOAD_IMAGE_FILE_SIZE_LIMIT", 15): with pytest.raises(FileTooLargeError): file_service.upload_file(filename="test.jpg", content=content, mimetype="image/jpeg", user=MagicMock()) def test_upload_file_end_user(self, file_service: FileService, mock_db_session): user = MagicMock(spec=EndUser) user.id = "end_user_id" with ( patch("services.file_service.storage"), patch("services.file_service.extract_tenant_id") as mock_tenant, patch("services.file_service.file_helpers.get_signed_file_url"), ): mock_tenant.return_value = "tenant" result = file_service.upload_file(filename="test.txt", content=b"test", mimetype="text/plain", user=user) assert result.created_by_role == CreatorUserRole.END_USER def test_is_file_size_within_limit(self): with ( patch.object(dify_config, "UPLOAD_IMAGE_FILE_SIZE_LIMIT", 10), patch.object(dify_config, "UPLOAD_VIDEO_FILE_SIZE_LIMIT", 20), patch.object(dify_config, "UPLOAD_AUDIO_FILE_SIZE_LIMIT", 30), patch.object(dify_config, "UPLOAD_FILE_SIZE_LIMIT", 5), ): # Image assert FileService.is_file_size_within_limit(extension="jpg", file_size=10 * 1024 * 1024) is True assert FileService.is_file_size_within_limit(extension="png", file_size=11 * 1024 * 1024) is False # Video assert FileService.is_file_size_within_limit(extension="mp4", file_size=20 * 1024 * 1024) is True assert FileService.is_file_size_within_limit(extension="avi", file_size=21 * 1024 * 1024) is False # Audio assert FileService.is_file_size_within_limit(extension="mp3", file_size=30 * 1024 * 1024) is True assert FileService.is_file_size_within_limit(extension="wav", file_size=31 * 1024 * 1024) is False # Default assert FileService.is_file_size_within_limit(extension="txt", file_size=5 * 1024 * 1024) is True assert FileService.is_file_size_within_limit(extension="pdf", file_size=6 * 1024 * 1024) is False def test_get_file_base64_success(self, file_service: FileService, mock_db_session): # Setup upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.key = "test_key" mock_db_session.scalar.return_value = upload_file with patch("services.file_service.storage") as mock_storage: mock_storage.load_once.return_value = b"test content" # Execute result = file_service.get_file_base64("file_id") # Assert assert result == base64.b64encode(b"test content").decode() mock_storage.load_once.assert_called_once_with("test_key") def test_get_file_base64_not_found(self, file_service: FileService, mock_db_session): mock_db_session.scalar.return_value = None with pytest.raises(NotFound, match="File not found"): file_service.get_file_base64("non_existent") def test_upload_text_success(self, file_service: FileService, mock_db_session): # Setup text = "sample text" text_name = "test.txt" user_id = "user_id" tenant_id = "tenant_id" with patch("services.file_service.storage") as mock_storage: # Execute result = file_service.upload_text(text, text_name, user_id, tenant_id) # Assert assert result.name == text_name assert result.size == len(text) assert result.tenant_id == tenant_id assert result.created_by == user_id assert result.used is True assert result.extension == "txt" mock_storage.save.assert_called_once() mock_db_session.add.assert_called_once() mock_db_session.commit.assert_called_once() def test_upload_text_long_name(self, file_service: FileService, mock_db_session): long_name = "a" * 210 with patch("services.file_service.storage"): result = file_service.upload_text("text", long_name, "user", "tenant") assert len(result.name) == 200 def test_get_file_preview_success(self, file_service: FileService, mock_db_session): # Setup upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.extension = "pdf" mock_db_session.scalar.return_value = upload_file with patch("services.file_service.ExtractProcessor.load_from_upload_file") as mock_extract: mock_extract.return_value = "Extracted text content" # Execute result = file_service.get_file_preview("file_id", "tenant_id") # Assert assert result == "Extracted text content" def test_get_file_preview_not_found(self, file_service: FileService, mock_db_session): mock_db_session.scalar.return_value = None with pytest.raises(NotFound, match="File not found"): file_service.get_file_preview("non_existent", "tenant_id") def test_get_file_preview_unsupported_type(self, file_service: FileService, mock_db_session): upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.extension = "exe" mock_db_session.scalar.return_value = upload_file with pytest.raises(UnsupportedFileTypeError): file_service.get_file_preview("file_id", "tenant_id") def test_get_image_preview_success(self, file_service: FileService, mock_db_session): # Setup upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.extension = "jpg" upload_file.mime_type = "image/jpeg" upload_file.key = "key" mock_db_session.scalar.return_value = upload_file with ( patch("services.file_service.file_helpers.verify_image_signature") as mock_verify, patch("services.file_service.storage") as mock_storage, ): mock_verify.return_value = True mock_storage.load.return_value = iter([b"chunk1"]) # Execute gen, mime = file_service.get_image_preview("file_id", "ts", "nonce", "sign") # Assert assert list(gen) == [b"chunk1"] assert mime == "image/jpeg" def test_get_image_preview_invalid_sig(self, file_service): with patch("services.file_service.file_helpers.verify_image_signature") as mock_verify: mock_verify.return_value = False with pytest.raises(NotFound, match="File not found or signature is invalid"): file_service.get_image_preview("file_id", "ts", "nonce", "sign") def test_get_image_preview_not_found(self, file_service: FileService, mock_db_session): mock_db_session.scalar.return_value = None with patch("services.file_service.file_helpers.verify_image_signature") as mock_verify: mock_verify.return_value = True with pytest.raises(NotFound, match="File not found or signature is invalid"): file_service.get_image_preview("file_id", "ts", "nonce", "sign") def test_get_image_preview_unsupported_type(self, file_service: FileService, mock_db_session): upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.extension = "txt" mock_db_session.scalar.return_value = upload_file with patch("services.file_service.file_helpers.verify_image_signature") as mock_verify: mock_verify.return_value = True with pytest.raises(UnsupportedFileTypeError): file_service.get_image_preview("file_id", "ts", "nonce", "sign") def test_get_file_generator_by_file_id_success(self, file_service: FileService, mock_db_session): upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.key = "key" mock_db_session.scalar.return_value = upload_file with ( patch("services.file_service.file_helpers.verify_file_signature") as mock_verify, patch("services.file_service.storage") as mock_storage, ): mock_verify.return_value = True mock_storage.load.return_value = iter([b"chunk"]) gen, file = file_service.get_file_generator_by_file_id("file_id", "ts", "nonce", "sign") assert list(gen) == [b"chunk"] assert file == upload_file def test_get_file_generator_by_file_id_invalid_sig(self, file_service): with patch("services.file_service.file_helpers.verify_file_signature") as mock_verify: mock_verify.return_value = False with pytest.raises(NotFound, match="File not found or signature is invalid"): file_service.get_file_generator_by_file_id("file_id", "ts", "nonce", "sign") def test_get_file_generator_by_file_id_not_found(self, file_service: FileService, mock_db_session): mock_db_session.scalar.return_value = None with patch("services.file_service.file_helpers.verify_file_signature") as mock_verify: mock_verify.return_value = True with pytest.raises(NotFound, match="File not found or signature is invalid"): file_service.get_file_generator_by_file_id("file_id", "ts", "nonce", "sign") def test_get_public_image_preview_success(self, file_service: FileService, mock_db_session): upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.extension = "png" upload_file.mime_type = "image/png" upload_file.key = "key" mock_db_session.scalar.return_value = upload_file with patch("services.file_service.storage") as mock_storage: mock_storage.load.return_value = b"image content" gen, mime = file_service.get_public_image_preview("file_id") assert gen == b"image content" assert mime == "image/png" def test_get_public_image_preview_not_found(self, file_service: FileService, mock_db_session): mock_db_session.scalar.return_value = None with pytest.raises(NotFound, match="File not found or signature is invalid"): file_service.get_public_image_preview("file_id") def test_get_public_image_preview_unsupported_type(self, file_service: FileService, mock_db_session): upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.extension = "txt" mock_db_session.scalar.return_value = upload_file with pytest.raises(UnsupportedFileTypeError): file_service.get_public_image_preview("file_id") def test_get_file_content_success(self, file_service: FileService, mock_db_session): upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.key = "key" mock_db_session.scalar.return_value = upload_file with patch("services.file_service.storage") as mock_storage: mock_storage.load.return_value = b"hello world" result = file_service.get_file_content("file_id") assert result == "hello world" def test_get_file_content_not_found(self, file_service: FileService, mock_db_session): mock_db_session.scalar.return_value = None with pytest.raises(NotFound, match="File not found"): file_service.get_file_content("file_id") def test_delete_file_success(self, file_service: FileService, mock_db_session): upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.key = "key" # For session.scalar(select(...)) mock_db_session.scalar.return_value = upload_file with patch("services.file_service.storage") as mock_storage: file_service.delete_file("file_id") mock_storage.delete.assert_called_once_with("key") mock_db_session.delete.assert_called_once_with(upload_file) def test_delete_file_not_found(self, file_service: FileService, mock_db_session): mock_db_session.scalar.return_value = None file_service.delete_file("file_id") # Should return without doing anything @patch("services.file_service.db") def test_get_upload_files_by_ids_empty(self, mock_db): result = FileService.get_upload_files_by_ids("tenant_id", []) assert result == {} @patch("services.file_service.db") def test_get_upload_files_by_ids(self, mock_db): upload_file = MagicMock(spec=UploadFile) upload_file.id = "550e8400-e29b-41d4-a716-446655440000" upload_file.tenant_id = "tenant_id" mock_db.session.scalars().all.return_value = [upload_file] result = FileService.get_upload_files_by_ids("tenant_id", ["550e8400-e29b-41d4-a716-446655440000"]) assert result["550e8400-e29b-41d4-a716-446655440000"] == upload_file def test_sanitize_zip_entry_name(self): assert FileService._sanitize_zip_entry_name("path/to/file.txt") == "file.txt" assert FileService._sanitize_zip_entry_name("../../../etc/passwd") == "passwd" assert FileService._sanitize_zip_entry_name(" ") == "file" assert FileService._sanitize_zip_entry_name("a\\b") == "a_b" def test_dedupe_zip_entry_name(self): used = {"a.txt"} assert FileService._dedupe_zip_entry_name("b.txt", used) == "b.txt" assert FileService._dedupe_zip_entry_name("a.txt", used) == "a (1).txt" used.add("a (1).txt") assert FileService._dedupe_zip_entry_name("a.txt", used) == "a (2).txt" def test_build_upload_files_zip_tempfile(self): upload_file = MagicMock(spec=UploadFile) upload_file.name = "test.txt" upload_file.key = "key" with ( patch("services.file_service.storage") as mock_storage, patch("services.file_service.os.remove") as mock_remove, ): mock_storage.load.return_value = [b"chunk1", b"chunk2"] with FileService.build_upload_files_zip_tempfile(upload_files=[upload_file]) as tmp_path: assert os.path.exists(tmp_path) mock_remove.assert_called_once()