140 lines
4.9 KiB
Python
140 lines
4.9 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
from fastapi import HTTPException
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from app.gateway.routers import system_settings as settings_router
|
|
from deerflow.persistence.knowledge_retrieval_settings import (
|
|
KnowledgeRetrievalSettingsRepository,
|
|
KnowledgeRetrievalSettingsRow,
|
|
MemoryKnowledgeRetrievalSettingsStore,
|
|
)
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def sql_store(tmp_path):
|
|
engine = create_async_engine(
|
|
f"sqlite+aiosqlite:///{tmp_path / 'knowledge-retrieval-settings.db'}"
|
|
)
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(KnowledgeRetrievalSettingsRow.__table__.create)
|
|
store = KnowledgeRetrievalSettingsRepository(
|
|
async_sessionmaker(engine, expire_on_commit=False)
|
|
)
|
|
try:
|
|
yield store
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sql_store_defaults_disabled_and_persists_updates(sql_store):
|
|
assert await sql_store.get() is False
|
|
assert await sql_store.get_advanced_management_enabled() is False
|
|
assert await sql_store.set(True, updated_by="admin-1") is True
|
|
assert await sql_store.get() is True
|
|
assert await sql_store.get_advanced_management_enabled() is False
|
|
assert (
|
|
await sql_store.set_advanced_management_enabled(
|
|
True,
|
|
updated_by="admin-1",
|
|
)
|
|
is True
|
|
)
|
|
assert await sql_store.get_advanced_management_enabled() is True
|
|
assert await sql_store.get() is True
|
|
assert await sql_store.set(False, updated_by="admin-1") is False
|
|
assert await sql_store.get() is False
|
|
assert await sql_store.get_advanced_management_enabled() is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_memory_store_defaults_disabled():
|
|
store = MemoryKnowledgeRetrievalSettingsStore()
|
|
assert await store.get() is False
|
|
assert await store.get_advanced_management_enabled() is False
|
|
assert await store.set(True) is True
|
|
assert await store.get() is True
|
|
assert await store.set_advanced_management_enabled(True) is True
|
|
assert await store.get_advanced_management_enabled() is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_api_reads_default_and_admin_can_update(monkeypatch):
|
|
store = MemoryKnowledgeRetrievalSettingsStore()
|
|
request = SimpleNamespace(
|
|
app=SimpleNamespace(
|
|
state=SimpleNamespace(knowledge_retrieval_settings_store=store)
|
|
)
|
|
)
|
|
|
|
async def admin(_request):
|
|
return SimpleNamespace(id="admin-1", system_role="admin")
|
|
|
|
monkeypatch.setattr(settings_router, "get_optional_user_from_request", admin)
|
|
initial = await settings_router.get_ordinary_knowledge_retrieval_setting(request)
|
|
assert initial.smart_reasoning_enabled is False
|
|
assert initial.advanced_management_enabled is False
|
|
|
|
updated = await settings_router.update_ordinary_knowledge_retrieval_setting(
|
|
settings_router.OrdinaryKnowledgeRetrievalUpdateRequest(
|
|
smart_reasoning_enabled=True
|
|
),
|
|
request,
|
|
)
|
|
assert updated.smart_reasoning_enabled is True
|
|
assert updated.advanced_management_enabled is False
|
|
|
|
advanced = await settings_router.update_ordinary_knowledge_retrieval_setting(
|
|
settings_router.OrdinaryKnowledgeRetrievalUpdateRequest(
|
|
advanced_management_enabled=True
|
|
),
|
|
request,
|
|
)
|
|
assert advanced.smart_reasoning_enabled is True
|
|
assert advanced.advanced_management_enabled is True
|
|
|
|
saved = await settings_router.get_ordinary_knowledge_retrieval_setting(request)
|
|
assert saved.smart_reasoning_enabled is True
|
|
assert saved.advanced_management_enabled is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_api_rejects_non_admin_without_changing_value(monkeypatch):
|
|
store = MemoryKnowledgeRetrievalSettingsStore()
|
|
request = SimpleNamespace(
|
|
app=SimpleNamespace(
|
|
state=SimpleNamespace(knowledge_retrieval_settings_store=store)
|
|
)
|
|
)
|
|
|
|
async def ordinary_user(_request):
|
|
return SimpleNamespace(id="user-1", system_role="user")
|
|
|
|
monkeypatch.setattr(
|
|
settings_router,
|
|
"get_optional_user_from_request",
|
|
ordinary_user,
|
|
)
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await settings_router.update_ordinary_knowledge_retrieval_setting(
|
|
settings_router.OrdinaryKnowledgeRetrievalUpdateRequest(
|
|
smart_reasoning_enabled=True
|
|
),
|
|
request,
|
|
)
|
|
assert exc_info.value.status_code == 403
|
|
assert await store.get() is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_api_reports_unavailable_store():
|
|
request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace()))
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await settings_router.get_ordinary_knowledge_retrieval_setting(request)
|
|
assert exc_info.value.status_code == 503
|