277 lines
9.1 KiB
Python
277 lines
9.1 KiB
Python
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_stage_copy_creates_an_independently_deletable_source_object(
|
||
tmp_path: Path,
|
||
) -> None:
|
||
storage = LocalDataProcessStorage(tmp_path / "storage")
|
||
original = _stage(storage, content=b"immutable source")
|
||
storage.publish([original])
|
||
|
||
copied = storage.stage_copy(
|
||
batch_id="batch-copy",
|
||
source_reference=original.reference,
|
||
expected_source_task_id="task-1",
|
||
expected_source_file_id="source-1",
|
||
task_id="task-2",
|
||
source_file_id="source-2",
|
||
version=1,
|
||
name="source.txt",
|
||
)
|
||
storage.publish([copied])
|
||
|
||
assert storage.read(copied.reference) == b"immutable source"
|
||
assert storage.delete(
|
||
original.reference,
|
||
expected_task_id="task-1",
|
||
expected_source_file_id="source-1",
|
||
) is True
|
||
assert storage.read(copied.reference) == b"immutable source"
|
||
assert storage.delete(
|
||
copied.reference,
|
||
expected_task_id="task-2",
|
||
expected_source_file_id="source-2",
|
||
) is True
|
||
_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()
|