YG_FT/backend/tests/test_data_process_storage.py

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()