from __future__ import annotations from pathlib import Path, PurePosixPath import pytest from app.modules.data_process import storage as storage_module from app.modules.data_process.storage import ( DataProcessStorageError, LocalDataProcessStorage, StagedSourceObject, ) def _stage( storage: LocalDataProcessStorage, *, batch_id: str = "batch-main", task_id: str = "task-1", source_file_id: str = "source-1", version: int = 1, name: str = "source.txt", content: bytes = b"payload", ) -> StagedSourceObject: return storage.stage_bytes( batch_id=batch_id, task_id=task_id, source_file_id=source_file_id, version=version, name=name, content=content, ) def _create_symlink(link: Path, target: Path, *, target_is_directory: bool = False) -> None: try: link.symlink_to(target, target_is_directory=target_is_directory) except (NotImplementedError, OSError) as exc: pytest.skip(f"当前平台不支持创建测试所需的符号链接: {exc}") def _assert_staging_empty(storage: LocalDataProcessStorage) -> None: assert list((storage.root / ".staging").iterdir()) == [] def test_stage_publish_read_delete_roundtrip_with_unicode_filename(tmp_path: Path) -> None: storage = LocalDataProcessStorage(tmp_path / "storage") content = "第一行\n第二行，100% 完成".encode() staged = _stage( storage, name="中文 数据 100%.csv", content=content, ) assert "%20" in staged.reference assert "%25" in staged.reference storage.publish([staged]) assert storage.read(staged.reference) == content assert storage.delete(staged.reference) is True assert storage.delete(staged.reference) is False _assert_staging_empty(storage) def test_db_reference_is_left_to_database_storage(tmp_path: Path) -> None: storage = LocalDataProcessStorage(tmp_path / "storage") assert storage.read("db://source-files/source-1") is None assert storage.delete("db://source-files/source-1") is False def test_owned_source_can_be_streamed_by_byte_range(tmp_path: Path) -> None: storage = LocalDataProcessStorage(tmp_path / "storage") content = b"0123456789abcdef" staged = _stage(storage, content=content) storage.publish([staged]) assert storage.file_size( staged.reference, expected_task_id="task-1", expected_source_file_id="source-1", ) == len(content) assert b"".join(storage.iter_bytes( staged.reference, expected_task_id="task-1", expected_source_file_id="source-1", expected_size=len(content), start=4, length=6, chunk_size=2, )) == b"456789" with pytest.raises(DataProcessStorageError, match="owner mismatch"): storage.file_size( staged.reference, expected_task_id="another-task", expected_source_file_id="source-1", ) with pytest.raises(DataProcessStorageError, match="does not match metadata"): b"".join(storage.iter_bytes( staged.reference, expected_task_id="task-1", expected_source_file_id="source-1", expected_size=len(content) + 1, )) @pytest.mark.parametrize( "reference", [ "local://data-process/../source-1/v1/file.txt", "local://data-process/task-1/source-1/v1/file%2Fname.txt", "local://data-process/task-1/source-1/v1/file.txt?download=1", "local://data-process/task-1/source-1/v1/file.txt#fragment", "https://data-process/task-1/source-1/v1/file.txt", ], ids=[ "parent-traversal", "percent-encoded-slash", "query", "fragment", "wrong-scheme", ], ) def test_unsafe_references_are_rejected(tmp_path: Path, reference: str) -> None: storage = LocalDataProcessStorage(tmp_path / "storage") with pytest.raises(DataProcessStorageError): storage.read(reference) with pytest.raises(DataProcessStorageError): storage.delete(reference) def test_publish_rejects_intermediate_directory_symlink(tmp_path: Path) -> None: storage = LocalDataProcessStorage(tmp_path / "storage") outside = tmp_path / "outside" outside.mkdir() staged = _stage(storage, task_id="linked-task") _create_symlink( storage.root / "linked-task", outside, target_is_directory=True, ) with pytest.raises(DataProcessStorageError, match="symlink|non-directory"): storage.publish([staged]) assert list(outside.iterdir()) == [] _assert_staging_empty(storage) def test_target_symlink_is_never_followed_or_deleted(tmp_path: Path) -> None: storage = LocalDataProcessStorage(tmp_path / "storage") staged = _stage(storage, task_id="task-link", source_file_id="source-link") outside_file = tmp_path / "outside.txt" outside_file.write_bytes(b"outside sentinel") final_path = storage.root.joinpath(*staged._relative_path.parts) final_path.parent.mkdir(parents=True) _create_symlink(final_path, outside_file) with pytest.raises(DataProcessStorageError, match="already exists"): storage.publish([staged]) with pytest.raises(DataProcessStorageError, match="regular file"): storage.read(staged.reference) with pytest.raises(DataProcessStorageError, match="non-regular"): storage.delete(staged.reference) assert final_path.is_symlink() assert outside_file.read_bytes() == b"outside sentinel" _assert_staging_empty(storage) def test_publish_rolls_back_first_object_when_second_target_collides(tmp_path: Path) -> None: storage = LocalDataProcessStorage(tmp_path / "storage") existing = _stage( storage, batch_id="batch-existing", source_file_id="source-existing", content=b"existing content", ) storage.publish([existing]) first = _stage( storage, batch_id="batch-new", source_file_id="source-new", content=b"must be rolled back", ) colliding_second = _stage( storage, batch_id="batch-new", source_file_id="source-existing", content=b"must not replace existing content", ) with pytest.raises(DataProcessStorageError, match="already exists"): storage.publish([first, colliding_second]) with pytest.raises(DataProcessStorageError, match="does not exist"): storage.read(first.reference) assert storage.read(existing.reference) == b"existing content" _assert_staging_empty(storage) def test_publish_rejects_manually_forged_staged_object(tmp_path: Path) -> None: storage = LocalDataProcessStorage(tmp_path / "storage") temporary_path = storage.root / ".staging" / "batch-forged" / "forged.tmp" temporary_path.parent.mkdir() temporary_path.write_bytes(b"forged content") relative_path = PurePosixPath("task-forged", "source-forged", "v1", "forged.txt") forged = StagedSourceObject( reference="local://data-process/task-forged/source-forged/v1/forged.txt", _temporary_path=temporary_path, _relative_path=relative_path, ) with pytest.raises(DataProcessStorageError, match="was not issued"): storage.publish([forged]) with pytest.raises(DataProcessStorageError, match="was not issued"): storage.discard([forged]) assert temporary_path.read_bytes() == b"forged content" assert not storage.root.joinpath(*relative_path.parts).exists() def test_relative_storage_configuration_is_anchored_to_backend_root( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: relative_configuration = Path("relative-storage") / tmp_path.name backend_root = Path(storage_module.__file__).resolve().parents[3] monkeypatch.chdir(tmp_path) monkeypatch.setenv("DATA_PROCESS_STORAGE_DIR", str(relative_configuration)) storage_module.get_data_process_storage.cache_clear() try: configured_root = storage_module._configured_storage_root() assert configured_root == backend_root / relative_configuration assert not configured_root.exists() finally: storage_module.get_data_process_storage.cache_clear()