154 lines
5.1 KiB
Python
154 lines
5.1 KiB
Python
from __future__ import annotations
|
|
|
|
from io import BytesIO
|
|
|
|
import pytest
|
|
|
|
from beaver.services.user_files import (
|
|
LocalUserFileStorage,
|
|
MinIOStorageConfig,
|
|
MinIOUserFileStorage,
|
|
UserFileNotFoundError,
|
|
UserFilePathError,
|
|
UserFileSizeError,
|
|
UserFileService,
|
|
normalize_user_path,
|
|
)
|
|
|
|
|
|
def test_normalize_user_path_accepts_fixed_roots() -> None:
|
|
assert normalize_user_path("uploads/readme.txt", allow_root=False) == "uploads/readme.txt"
|
|
assert normalize_user_path("outputs/report.md", allow_root=False) == "outputs/report.md"
|
|
assert normalize_user_path("tasks/task-123/draft.md", allow_root=False) == "tasks/task-123/draft.md"
|
|
assert normalize_user_path("", allow_root=True) == ""
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"path",
|
|
[
|
|
"../secret.txt",
|
|
"/uploads/input.txt",
|
|
"/outputs/result.txt",
|
|
"/shared/profile.json",
|
|
"/tasks/task-123/draft.md",
|
|
"uploads/../state/config.json",
|
|
"memory/private.txt",
|
|
"uploads/.internal",
|
|
"",
|
|
],
|
|
)
|
|
def test_normalize_user_path_rejects_invalid_paths(path: str) -> None:
|
|
with pytest.raises(UserFilePathError):
|
|
normalize_user_path(path, allow_root=False)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_file_service_root_and_round_trip(tmp_path) -> None:
|
|
service = UserFileService(LocalUserFileStorage(tmp_path / "user-files"))
|
|
|
|
root = await service.browse("")
|
|
uploaded = await service.upload(
|
|
"uploads",
|
|
"hello.txt",
|
|
b"hello user files",
|
|
content_type="text/plain",
|
|
)
|
|
uploads = await service.browse("uploads")
|
|
preview = await service.preview("uploads/hello.txt")
|
|
downloaded = await service.download("uploads/hello.txt")
|
|
deleted = await service.delete("uploads/hello.txt")
|
|
|
|
assert [item["name"] for item in root["items"]] == ["uploads", "outputs", "shared", "tasks"]
|
|
assert uploaded["path"] == "uploads/hello.txt"
|
|
assert uploaded["content_type"] == "text/plain"
|
|
assert [item["name"] for item in uploads["items"]] == ["hello.txt"]
|
|
assert preview["content"] == "hello user files"
|
|
assert downloaded.content == b"hello user files"
|
|
assert deleted is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_file_service_stream_upload_and_size_limit(tmp_path) -> None:
|
|
service = UserFileService(LocalUserFileStorage(tmp_path / "user-files"))
|
|
|
|
uploaded = await service.upload_stream(
|
|
"uploads",
|
|
"streamed.txt",
|
|
BytesIO(b"streamed user file"),
|
|
content_type="text/plain",
|
|
max_bytes=1024,
|
|
part_size=4,
|
|
)
|
|
preview = await service.preview("uploads/streamed.txt")
|
|
|
|
assert uploaded["path"] == "uploads/streamed.txt"
|
|
assert uploaded["size"] == len(b"streamed user file")
|
|
assert preview["content"] == "streamed user file"
|
|
|
|
with pytest.raises(UserFileSizeError):
|
|
await service.upload_stream(
|
|
"uploads",
|
|
"too-large.txt",
|
|
BytesIO(b"abcdef"),
|
|
content_type="text/plain",
|
|
max_bytes=5,
|
|
part_size=2,
|
|
)
|
|
with pytest.raises(UserFileNotFoundError):
|
|
await service.preview("uploads/too-large.txt")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_file_service_rejects_root_delete_and_traversal(tmp_path) -> None:
|
|
service = UserFileService(LocalUserFileStorage(tmp_path / "user-files"))
|
|
|
|
with pytest.raises(UserFilePathError):
|
|
await service.delete("uploads")
|
|
|
|
with pytest.raises(UserFilePathError):
|
|
await service.upload("../workspace", "hello.txt", b"x", content_type="text/plain")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_file_service_creates_nested_directories(tmp_path) -> None:
|
|
service = UserFileService(LocalUserFileStorage(tmp_path / "user-files"))
|
|
|
|
created = await service.mkdir("tasks/task-123/references")
|
|
tasks = await service.browse("tasks/task-123")
|
|
|
|
assert created["path"] == "tasks/task-123/references"
|
|
assert created["type"] == "directory"
|
|
assert [item["name"] for item in tasks["items"]] == ["references"]
|
|
|
|
|
|
def test_minio_storage_maps_virtual_paths_under_namespace() -> None:
|
|
storage = object.__new__(MinIOUserFileStorage)
|
|
storage.config = MinIOStorageConfig(
|
|
endpoint="minio.local:9000",
|
|
access_key="alice-access",
|
|
secret_key="alice-secret",
|
|
bucket="beaver-user-files",
|
|
namespace="users/alice",
|
|
)
|
|
|
|
assert storage._object_name("uploads/report.pdf") == "users/alice/uploads/report.pdf"
|
|
assert storage._object_name("tasks/task-123/result.json") == "users/alice/tasks/task-123/result.json"
|
|
assert storage._user_path("users/alice/outputs/summary.md") == "outputs/summary.md"
|
|
|
|
|
|
def test_minio_storage_rejects_paths_that_escape_namespace() -> None:
|
|
storage = object.__new__(MinIOUserFileStorage)
|
|
storage.config = MinIOStorageConfig(
|
|
endpoint="minio.local:9000",
|
|
access_key="alice-access",
|
|
secret_key="alice-secret",
|
|
bucket="beaver-user-files",
|
|
namespace="users/alice",
|
|
)
|
|
|
|
with pytest.raises(UserFilePathError):
|
|
storage._object_name("uploads/../state/config.json")
|
|
|
|
with pytest.raises(UserFilePathError):
|
|
storage._user_path("users/bob/uploads/secret.txt")
|