140 lines
5.3 KiB
Python
140 lines
5.3 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
from fastapi import HTTPException
|
|
from pydantic import ValidationError
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from app.gateway.routers import upload_policy as policy_router
|
|
from deerflow.persistence.upload_policy import MemoryUploadPolicyStore, UploadPolicyRepository, UploadPolicyRow
|
|
from deerflow.uploads.policy import MIB, default_upload_policy, get_runtime_large_file_policy
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def sql_store(tmp_path):
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'upload-policy.db'}")
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(UploadPolicyRow.__table__.create)
|
|
store = UploadPolicyRepository(async_sessionmaker(engine, expire_on_commit=False))
|
|
try:
|
|
yield store
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
def test_default_policy_matches_product_limits():
|
|
policy = default_upload_policy()
|
|
assert policy["max_files_per_message"] == 20
|
|
assert policy["max_total_bytes"] == 100 * MIB
|
|
assert policy["groups"]["document"] == {
|
|
"max_bytes": 10 * MIB,
|
|
"extensions": [".txt", ".md", ".doc", ".docx", ".ppt", ".pptx", ".pdf"],
|
|
}
|
|
assert policy["groups"]["image"]["max_bytes"] == 10 * MIB
|
|
assert policy["groups"]["spreadsheet"]["max_bytes"] == 5 * MIB
|
|
assert policy["large_file"] == {
|
|
"threshold_bytes": MIB,
|
|
"max_result_chars": 12_000,
|
|
"hard_max_result_chars": 32_000,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_new_sql_store_is_seeded_and_seed_is_idempotent(sql_store: UploadPolicyRepository):
|
|
first = await sql_store.ensure_seed(default_upload_policy())
|
|
assert first["policy"]["max_files_per_message"] == 20
|
|
assert first["revision"] == 1
|
|
|
|
changed = default_upload_policy()
|
|
changed["max_files_per_message"] = 7
|
|
saved = await sql_store.set_policy(changed, updated_by="admin-1")
|
|
assert saved["revision"] == 2
|
|
|
|
second_seed = await sql_store.ensure_seed(default_upload_policy())
|
|
assert second_seed["policy"]["max_files_per_message"] == 7
|
|
assert second_seed["revision"] == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_memory_store_has_sql_parity_and_revision():
|
|
store = MemoryUploadPolicyStore()
|
|
seeded = await store.ensure_seed(default_upload_policy())
|
|
assert seeded["revision"] == 1
|
|
policy = default_upload_policy()
|
|
policy["large_file"]["max_result_chars"] = 9000
|
|
saved = await store.set_policy(policy, updated_by="admin")
|
|
assert saved["revision"] == 2
|
|
assert saved["updated_by"] == "admin"
|
|
|
|
|
|
def test_policy_rejects_overlap_unsupported_extensions_and_unsafe_context_limits():
|
|
policy = default_upload_policy()
|
|
policy["groups"]["image"]["extensions"].append(".txt")
|
|
with pytest.raises(ValidationError):
|
|
policy_router.UploadPolicyDocument.model_validate(policy)
|
|
|
|
policy = default_upload_policy()
|
|
policy["groups"]["document"]["extensions"].append(".exe")
|
|
with pytest.raises(ValidationError):
|
|
policy_router.UploadPolicyDocument.model_validate(policy)
|
|
|
|
policy = default_upload_policy()
|
|
policy["large_file"]["max_result_chars"] = 20_000
|
|
policy["large_file"]["hard_max_result_chars"] = 10_000
|
|
with pytest.raises(ValidationError):
|
|
policy_router.UploadPolicyDocument.model_validate(policy)
|
|
|
|
|
|
def test_per_file_config_hard_cap_matches_image_reader_cap():
|
|
policy = default_upload_policy()
|
|
policy["max_total_bytes"] = 200 * MIB
|
|
policy["groups"]["image"]["max_bytes"] = 100 * MIB
|
|
assert policy_router.UploadPolicyDocument.model_validate(policy).groups.image.max_bytes == 100 * MIB
|
|
|
|
policy["groups"]["image"]["max_bytes"] += 1
|
|
with pytest.raises(ValidationError):
|
|
policy_router.UploadPolicyDocument.model_validate(policy)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_put_updates_runtime_registry_immediately(monkeypatch):
|
|
store = MemoryUploadPolicyStore()
|
|
await store.ensure_seed(default_upload_policy())
|
|
request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(upload_policy_store=store)))
|
|
admin = SimpleNamespace(id="admin-1", system_role="admin")
|
|
|
|
async def actor(_request):
|
|
return admin
|
|
|
|
monkeypatch.setattr(policy_router, "_actor", actor)
|
|
policy = default_upload_policy()
|
|
policy["large_file"]["threshold_bytes"] = 2 * MIB
|
|
policy["large_file"]["max_result_chars"] = 8192
|
|
response = await policy_router.replace_upload_policy(
|
|
request,
|
|
policy_router.UploadPolicyDocument.model_validate(policy),
|
|
)
|
|
assert response.revision == 2
|
|
assert response.updated_by == "admin-1"
|
|
assert get_runtime_large_file_policy()["threshold_bytes"] == 2 * MIB
|
|
assert get_runtime_large_file_policy()["max_result_chars"] == 8192
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_put_rejects_non_admin(monkeypatch):
|
|
request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(upload_policy_store=MemoryUploadPolicyStore())))
|
|
|
|
async def actor(_request):
|
|
return SimpleNamespace(id="user-1", system_role="user")
|
|
|
|
monkeypatch.setattr(policy_router, "_actor", actor)
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await policy_router.replace_upload_policy(
|
|
request,
|
|
policy_router.UploadPolicyDocument.model_validate(default_upload_policy()),
|
|
)
|
|
assert exc_info.value.status_code == 403
|