143 lines
4.7 KiB
Python
143 lines
4.7 KiB
Python
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()
|