geointel/backend/tests/test_storage_service.py

import asyncio
from hashlib import sha256
from pathlib import Path
from types import SimpleNamespace

import pytest

from app.core.errors import AppError
from app.services.dataset_service import DatasetService
from app.services.storage_service import StorageService


def test_persist_dataset_file_records_metadata(monkeypatch, tmp_path) -> None:
    monkeypatch.setattr(
        "app.services.storage_service.get_settings",
        lambda: SimpleNamespace(storage_root=str(tmp_path)),
    )

    metadata = StorageService.persist_dataset_file(
        project_id="project-123",
        dataset_id="dataset-456",
        dataset_type="vector",
        original_filename="../weird name!@#.geojson",
        content=b"example-bytes",
        content_type="application/geo+json",
    )

    assert metadata["original_filename"] == "weird name___.geojson"
    assert metadata["stored_filename"] == "dataset-456_weird name___.geojson"
    assert metadata["content_type"] == "application/geo+json"
    assert metadata["size_bytes"] == 13
    assert len(metadata["checksum_sha256"]) == 64
    assert Path(metadata["storage_path"]).exists()
    assert str(Path(tmp_path, "uploads", "project-123", "vector", "dataset-456")) in metadata["storage_path"]


class _ChunkedUpload:
    def __init__(self, content: bytes) -> None:
        self.content = content
        self.requested_sizes: list[int] = []

    async def read(self, size: int) -> bytes:
        self.requested_sizes.append(size)
        chunk, self.content = self.content[:size], self.content[size:]
        return chunk


def test_persist_upload_file_streams_bounded_chunks(monkeypatch, tmp_path) -> None:
    monkeypatch.setattr(StorageService, "_base_dir", staticmethod(lambda: tmp_path))
    upload = _ChunkedUpload(b"abcdefghijk")

    metadata = asyncio.run(
        StorageService.persist_upload_file(
            project_id="project",
            dataset_id="dataset",
            dataset_type="raster",
            original_filename="source.tif",
            upload=upload,
            content_type="image/tiff",
            max_bytes=32,
            chunk_size=4,
        )
    )

    stored = Path(metadata["storage_path"])
    assert upload.requested_sizes == [4, 4, 4, 4]
    assert stored.read_bytes() == b"abcdefghijk"
    assert metadata["size_bytes"] == 11
    assert metadata["checksum_sha256"] == sha256(b"abcdefghijk").hexdigest()


def test_persist_upload_file_rejects_oversize_and_removes_partial_file(monkeypatch, tmp_path) -> None:
    monkeypatch.setattr(StorageService, "_base_dir", staticmethod(lambda: tmp_path))
    upload = _ChunkedUpload(b"0123456789")
    expected_path = Path(
        StorageService.dataset_file_path(
            "project",
            "dataset",
            "vector",
            "source.geojson",
        )
    )

    with pytest.raises(AppError) as exc_info:
        asyncio.run(
            StorageService.persist_upload_file(
                project_id="project",
                dataset_id="dataset",
                dataset_type="vector",
                original_filename="source.geojson",
                upload=upload,
                content_type="application/geo+json",
                max_bytes=8,
                chunk_size=3,
            )
        )

    assert exc_info.value.code == "UPLOAD_TOO_LARGE"
    assert exc_info.value.status_code == 413
    assert not expected_path.exists()


def test_vector_staging_uses_lower_in_memory_limit(monkeypatch) -> None:
    captured = {}

    async def persist_upload_file(**kwargs):
        captured.update(kwargs)
        return {"storage_path": "unused"}

    monkeypatch.setattr(
        "app.services.dataset_service.get_settings",
        lambda: SimpleNamespace(max_upload_mb=500, max_in_memory_vector_mb=32),
    )
    monkeypatch.setattr(StorageService, "persist_upload_file", persist_upload_file)

    asyncio.run(
        DatasetService._stage_upload(
            project_id="project",
            dataset_id="dataset",
            dataset_type="vector",
            filename="source.geojson",
            file=SimpleNamespace(content_type="application/geo+json"),
        )
    )

    assert captured["max_bytes"] == 32 * 1024 * 1024


def test_staged_vector_read_is_bounded_before_loading_file(monkeypatch, tmp_path) -> None:
    source = tmp_path / "large.geojson"
    source.write_bytes(b"x" * (1024 * 1024 + 1))
    monkeypatch.setattr(
        "app.services.dataset_service.get_settings",
        lambda: SimpleNamespace(max_upload_mb=500, max_in_memory_vector_mb=1),
    )

    with pytest.raises(AppError) as exc_info:
        DatasetService._read_staged_vector_bytes({"storage_path": str(source)})

    assert exc_info.value.code == "UPLOAD_TOO_LARGE"
    assert exc_info.value.status_code == 413
    assert not source.exists()