177 lines
6.5 KiB
Python
177 lines
6.5 KiB
Python
"""资料导入候选抽取测试（5b，异步后台 + mock ai_provider）。"""
|
||
|
||
from collections.abc import AsyncGenerator
|
||
|
||
import pytest
|
||
from httpx import ASGITransport, AsyncClient
|
||
from sqlalchemy import select
|
||
|
||
from app.application.imports import drain_pending_extractions
|
||
from app.db.session import AsyncSessionLocal
|
||
from app.main import app
|
||
from app.models.user import User
|
||
from app.services.auth import create_access_token
|
||
|
||
|
||
async def _get_character_id(client: AsyncClient) -> str:
|
||
"""通过已认证的 API 获取当前用户活跃角色 ID。"""
|
||
resp = await client.get("/api/v1/characters/active")
|
||
assert resp.status_code == 200
|
||
return resp.json()["id"]
|
||
|
||
|
||
class _FakeChatResult:
|
||
def __init__(self, content: str) -> None:
|
||
self.content = content
|
||
|
||
|
||
_PROFILE_RESPONSE = (
|
||
'[{"field":"background","value":"来自北京，从事设计工作",'
|
||
'"confidence":0.9,"reason":"资料明确提及","evidence":[0,1]}]'
|
||
)
|
||
_MEMORY_RESPONSE = (
|
||
'[{"type":"preference","title":"喜欢简洁回答","content":"用户偏好简洁直接的回答",'
|
||
'"importance":0.8,"sensitivity":"low","reason":"多次提及","evidence":[0]}]'
|
||
)
|
||
_PERSONALITY_RESPONSE = (
|
||
'[{"field":"expression_style","change":"温和且有条理","reason":"资料体现","evidence":[1]}]'
|
||
)
|
||
|
||
|
||
@pytest.fixture
|
||
async def client(
|
||
monkeypatch: pytest.MonkeyPatch, auth_headers: dict[str, str]
|
||
) -> AsyncGenerator[AsyncClient, None]:
|
||
async def fake_chat(messages, *, model_hint, **_kwargs):
|
||
if model_hint == "memory-extraction":
|
||
return _FakeChatResult(_MEMORY_RESPONSE)
|
||
system_text = str(messages[0].content) if messages else ""
|
||
if "人格" in system_text:
|
||
return _FakeChatResult(_PERSONALITY_RESPONSE)
|
||
return _FakeChatResult(_PROFILE_RESPONSE)
|
||
|
||
class _MockMessage:
|
||
def __init__(self, content: str):
|
||
self.content = content
|
||
|
||
class _MockManager:
|
||
async def chat(self, messages, model_hint=None):
|
||
return await fake_chat(messages, model_hint=model_hint)
|
||
|
||
def to_message_list(self, system, _messages):
|
||
return [_MockMessage(s) for s in system]
|
||
|
||
monkeypatch.setattr("app.application.imports._get_provider_manager", lambda: _MockManager())
|
||
transport = ASGITransport(app=app)
|
||
async with AsyncClient(
|
||
transport=transport, base_url="http://test", headers=auth_headers
|
||
) as http_client:
|
||
yield http_client
|
||
|
||
|
||
async def _create_import(client: AsyncClient, character_id: str) -> str:
|
||
resp = await client.post(
|
||
f"/api/v1/characters/{character_id}/imports",
|
||
json={
|
||
"material_type": "background_story",
|
||
"content": "我叫小魂，来自北京。\n\n我从事设计工作，喜欢简洁直接的回答。",
|
||
},
|
||
)
|
||
return resp.json()["id"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_extract_profile_candidates(client: AsyncClient) -> None:
|
||
character_id = await _get_character_id(client)
|
||
import_id = await _create_import(client, character_id)
|
||
|
||
triggered = await client.post(
|
||
f"/api/v1/characters/{character_id}/imports/{import_id}/profile-candidates"
|
||
)
|
||
assert triggered.status_code == 200
|
||
assert triggered.json()["status"] == "extracting_profile"
|
||
|
||
await drain_pending_extractions()
|
||
|
||
candidates = (
|
||
await client.get(
|
||
f"/api/v1/characters/{character_id}/imports/{import_id}/profile-candidates"
|
||
)
|
||
).json()
|
||
assert len(candidates) == 1
|
||
assert candidates[0]["field"] == "background"
|
||
assert candidates[0]["status"] == "pending"
|
||
assert len(candidates[0]["evidence_ids"]) == 2 # 映射到 import_items.id
|
||
|
||
job = (await client.get(f"/api/v1/characters/{character_id}/imports/{import_id}")).json()
|
||
assert job["status"] == "reviewing"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_extract_memory_candidates(client: AsyncClient) -> None:
|
||
character_id = await _get_character_id(client)
|
||
import_id = await _create_import(client, character_id)
|
||
|
||
await client.post(f"/api/v1/characters/{character_id}/imports/{import_id}/memory-candidates")
|
||
await drain_pending_extractions()
|
||
|
||
candidates = (
|
||
await client.get(f"/api/v1/characters/{character_id}/imports/{import_id}/memory-candidates")
|
||
).json()
|
||
assert len(candidates) == 1
|
||
assert candidates[0]["type"] == "preference"
|
||
assert candidates[0]["title"] == "喜欢简洁回答"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_extract_personality_proposals(client: AsyncClient) -> None:
|
||
character_id = await _get_character_id(client)
|
||
import_id = await _create_import(client, character_id)
|
||
|
||
await client.post(
|
||
f"/api/v1/characters/{character_id}/imports/{import_id}/personality-proposals"
|
||
)
|
||
await drain_pending_extractions()
|
||
|
||
# 人格建议落入 personality_proposals（pending）；无独立列表端点，以 job 进入 reviewing 验证
|
||
job = (await client.get(f"/api/v1/characters/{character_id}/imports/{import_id}")).json()
|
||
assert job["status"] == "reviewing"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_extraction_failure_marks_job_failed(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
async def failing_chat(*_args, **_kwargs):
|
||
raise RuntimeError("模型不可用")
|
||
|
||
class _MockMessage:
|
||
def __init__(self, content: str):
|
||
self.content = content
|
||
|
||
class _FailingManager:
|
||
async def chat(self, messages, model_hint=None):
|
||
return await failing_chat(messages, model_hint=model_hint)
|
||
|
||
def to_message_list(self, system, _messages):
|
||
return [_MockMessage(s) for s in system]
|
||
|
||
monkeypatch.setattr("app.application.imports._get_provider_manager", lambda: _FailingManager())
|
||
transport = ASGITransport(app=app)
|
||
async with AsyncSessionLocal() as session:
|
||
result = await session.execute(select(User).where(User.is_admin.is_(True)).limit(1))
|
||
admin = result.scalar_one_or_none()
|
||
assert admin is not None
|
||
token = create_access_token(subject=admin.id, extra={"owner_scope": admin.id})
|
||
headers = {"Authorization": f"Bearer {token}"}
|
||
async with AsyncClient(transport=transport, base_url="http://test", headers=headers) as client:
|
||
character_id = await _get_character_id(client)
|
||
import_id = await _create_import(client, character_id)
|
||
|
||
await client.post(
|
||
f"/api/v1/characters/{character_id}/imports/{import_id}/profile-candidates"
|
||
)
|
||
await drain_pending_extractions()
|
||
|
||
job = (await client.get(f"/api/v1/characters/{character_id}/imports/{import_id}")).json()
|
||
assert job["status"] == "failed"
|
||
assert job["error_code"] == "extraction_failed"
|