339 lines
11 KiB
Python
339 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
from contextlib import asynccontextmanager
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
import yuxi.services.agent_run_service as agent_run_service
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_agent_run_events_emits_error_on_db_error(monkeypatch: pytest.MonkeyPatch):
|
|
@asynccontextmanager
|
|
async def fake_session_ctx():
|
|
yield object()
|
|
|
|
class BrokenRepo:
|
|
def __init__(self, db):
|
|
self.db = db
|
|
|
|
async def get_run_for_user(self, run_id: str, uid: str):
|
|
del run_id, uid
|
|
raise RuntimeError("db down")
|
|
|
|
monkeypatch.setattr(agent_run_service.pg_manager, "get_async_session_context", fake_session_ctx)
|
|
monkeypatch.setattr(agent_run_service, "AgentRunRepository", BrokenRepo)
|
|
|
|
chunks = []
|
|
async for chunk in agent_run_service.stream_agent_run_events(
|
|
run_id="run-1",
|
|
after_seq="0",
|
|
current_uid="user-1",
|
|
):
|
|
chunks.append(chunk)
|
|
|
|
assert len(chunks) == 1
|
|
assert chunks[0].startswith("event: error")
|
|
assert '"reason": "db_error"' in chunks[0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_agent_run_events_reads_redis_and_ends_on_end_event(monkeypatch: pytest.MonkeyPatch):
|
|
@asynccontextmanager
|
|
async def fake_session_ctx():
|
|
yield object()
|
|
|
|
class Repo:
|
|
def __init__(self, db):
|
|
self.db = db
|
|
|
|
async def get_run_for_user(self, run_id: str, uid: str):
|
|
del run_id, uid
|
|
return SimpleNamespace(status="completed", thread_id="thread-1")
|
|
|
|
calls = {"count": 0}
|
|
|
|
async def fake_list_events(run_id: str, *, after_seq: str, limit: int):
|
|
del run_id, after_seq, limit
|
|
calls["count"] += 1
|
|
if calls["count"] == 1:
|
|
return [
|
|
{
|
|
"seq": "1700000000000-0",
|
|
"event_type": "messages",
|
|
"payload": {
|
|
"schema_version": 1,
|
|
"run_id": "run-1",
|
|
"thread_id": "thread-1",
|
|
"event": "messages",
|
|
"payload": {"items": [{"status": "loading", "response": "你"}]},
|
|
"created_at": "2026-05-27T00:00:00+00:00",
|
|
},
|
|
"ts": 1700000000000,
|
|
},
|
|
{
|
|
"seq": "1700000000001-0",
|
|
"event_type": "end",
|
|
"payload": {
|
|
"schema_version": 1,
|
|
"run_id": "run-1",
|
|
"thread_id": "thread-1",
|
|
"event": "end",
|
|
"payload": {"status": "completed"},
|
|
"created_at": "2026-05-27T00:00:01+00:00",
|
|
},
|
|
"ts": 1700000000001,
|
|
},
|
|
]
|
|
return []
|
|
|
|
monkeypatch.setattr(agent_run_service.pg_manager, "get_async_session_context", fake_session_ctx)
|
|
monkeypatch.setattr(agent_run_service, "AgentRunRepository", Repo)
|
|
monkeypatch.setattr(agent_run_service, "list_run_stream_events", fake_list_events)
|
|
monkeypatch.setattr(agent_run_service, "SSE_POLL_INTERVAL_SECONDS", 0)
|
|
|
|
chunks = []
|
|
async for chunk in agent_run_service.stream_agent_run_events(
|
|
run_id="run-1",
|
|
after_seq="0",
|
|
current_uid="user-1",
|
|
):
|
|
chunks.append(chunk)
|
|
|
|
assert chunks[0].startswith("event: messages")
|
|
assert "id: 1700000000000-0" in chunks[0]
|
|
assert chunks[-1].startswith("event: end")
|
|
assert "id: 1700000000001-0" in chunks[-1]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_agent_run_persists_input_before_enqueue(monkeypatch: pytest.MonkeyPatch):
|
|
class FakeResult:
|
|
def scalar_one_or_none(self):
|
|
return SimpleNamespace(uid="user-1", role="user")
|
|
|
|
class FakeDB:
|
|
def __init__(self):
|
|
self.order: list[str] = []
|
|
self.committed = False
|
|
self.added = []
|
|
|
|
async def execute(self, stmt):
|
|
del stmt
|
|
return FakeResult()
|
|
|
|
def add(self, item):
|
|
self.added.append(item)
|
|
|
|
async def flush(self):
|
|
self.order.append("flush")
|
|
for item in self.added:
|
|
if getattr(item, "id", None) is None:
|
|
item.id = 10
|
|
|
|
async def commit(self):
|
|
self.order.append("commit")
|
|
self.committed = True
|
|
|
|
async def rollback(self):
|
|
raise AssertionError("rollback should not be called")
|
|
|
|
db = FakeDB()
|
|
created_run = SimpleNamespace(
|
|
id="",
|
|
thread_id="thread-1",
|
|
status="pending",
|
|
request_id="req-1",
|
|
uid="user-1",
|
|
)
|
|
|
|
class RunRepo:
|
|
def __init__(self, db_session):
|
|
self.db = db_session
|
|
|
|
async def get_run_by_request_id(self, request_id: str):
|
|
del request_id
|
|
return None
|
|
|
|
async def create_run(self, **kwargs):
|
|
assert kwargs["request_id"] == "req-1"
|
|
assert kwargs["conversation_id"] == 1
|
|
created_run.id = kwargs["run_id"]
|
|
return created_run
|
|
|
|
async def set_input_message(self, run_id: str, message_id: int):
|
|
assert run_id == created_run.id
|
|
assert message_id == 10
|
|
return created_run
|
|
|
|
class ConvRepo:
|
|
def __init__(self, db_session):
|
|
self.db = db_session
|
|
|
|
async def get_conversation_by_thread_id(self, thread_id: str):
|
|
del thread_id
|
|
return SimpleNamespace(id=1, uid="user-1", status="active", agent_id="default")
|
|
|
|
class AgentRepo:
|
|
def __init__(self, db_session):
|
|
self.db = db_session
|
|
|
|
async def get_visible_by_slug(self, slug: str, user):
|
|
del user
|
|
return SimpleNamespace(slug=slug, backend_id="ChatbotAgent")
|
|
|
|
class Queue:
|
|
async def enqueue_job(self, job_name: str, run_id: str, _job_id: str):
|
|
assert job_name == "process_agent_run"
|
|
assert run_id == created_run.id
|
|
assert _job_id == f"run:{created_run.id}"
|
|
db.order.append("enqueue")
|
|
assert db.committed is True
|
|
|
|
async def fake_get_arq_pool():
|
|
return Queue()
|
|
|
|
monkeypatch.setattr(agent_run_service.agent_manager, "get_agent", lambda backend_id: object())
|
|
monkeypatch.setattr(agent_run_service, "AgentRepository", AgentRepo)
|
|
monkeypatch.setattr(agent_run_service, "ConversationRepository", ConvRepo)
|
|
monkeypatch.setattr(agent_run_service, "AgentRunRepository", RunRepo)
|
|
monkeypatch.setattr(agent_run_service, "get_arq_pool", fake_get_arq_pool)
|
|
|
|
result = await agent_run_service.create_agent_run_view(
|
|
query="hello",
|
|
agent_id="default",
|
|
thread_id="thread-1",
|
|
meta={"request_id": "req-1"},
|
|
image_content=None,
|
|
current_uid="user-1",
|
|
db=db,
|
|
)
|
|
|
|
assert db.order[-2:] == ["commit", "enqueue"]
|
|
assert result["run_id"] == created_run.id
|
|
assert result["request_id"] == "req-1"
|
|
assert db.added[0].run_id == created_run.id
|
|
assert db.added[0].request_id == "req-1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_resume_run_marks_input_message_source(monkeypatch: pytest.MonkeyPatch):
|
|
class FakeResult:
|
|
def scalar_one_or_none(self):
|
|
return SimpleNamespace(uid="user-1", role="user")
|
|
|
|
class FakeDB:
|
|
def __init__(self):
|
|
self.order: list[str] = []
|
|
self.committed = False
|
|
self.added = []
|
|
|
|
async def execute(self, stmt):
|
|
del stmt
|
|
return FakeResult()
|
|
|
|
def add(self, item):
|
|
self.added.append(item)
|
|
|
|
async def flush(self):
|
|
self.order.append("flush")
|
|
for item in self.added:
|
|
if getattr(item, "id", None) is None:
|
|
item.id = 11
|
|
|
|
async def commit(self):
|
|
self.order.append("commit")
|
|
self.committed = True
|
|
|
|
async def rollback(self):
|
|
raise AssertionError("rollback should not be called")
|
|
|
|
db = FakeDB()
|
|
created_run = SimpleNamespace(
|
|
id="",
|
|
thread_id="thread-1",
|
|
status="pending",
|
|
request_id="resume-req",
|
|
uid="user-1",
|
|
)
|
|
|
|
class RunRepo:
|
|
def __init__(self, db_session):
|
|
self.db = db_session
|
|
|
|
async def get_run_for_user(self, run_id: str, uid: str):
|
|
assert run_id == "parent-run"
|
|
assert uid == "user-1"
|
|
return SimpleNamespace(id=run_id, thread_id="thread-1", status="interrupted")
|
|
|
|
async def get_resume_run(self, parent_run_id: str, resume_request_id: str):
|
|
assert parent_run_id == "parent-run"
|
|
assert resume_request_id == "resume-req"
|
|
return None
|
|
|
|
async def get_run_by_request_id(self, request_id: str):
|
|
assert request_id == "resume-req"
|
|
return None
|
|
|
|
async def create_run(self, **kwargs):
|
|
assert kwargs["run_type"] == "resume"
|
|
assert kwargs["parent_run_id"] == "parent-run"
|
|
assert kwargs["resume_request_id"] == "resume-req"
|
|
created_run.id = kwargs["run_id"]
|
|
return created_run
|
|
|
|
async def set_input_message(self, run_id: str, message_id: int):
|
|
assert run_id == created_run.id
|
|
assert message_id == 11
|
|
return created_run
|
|
|
|
class ConvRepo:
|
|
def __init__(self, db_session):
|
|
self.db = db_session
|
|
|
|
async def get_conversation_by_thread_id(self, thread_id: str):
|
|
del thread_id
|
|
return SimpleNamespace(id=1, uid="user-1", status="active", agent_id="default")
|
|
|
|
class AgentRepo:
|
|
def __init__(self, db_session):
|
|
self.db = db_session
|
|
|
|
async def get_visible_by_slug(self, slug: str, user):
|
|
del user
|
|
return SimpleNamespace(slug=slug, backend_id="ChatbotAgent")
|
|
|
|
class Queue:
|
|
async def enqueue_job(self, job_name: str, run_id: str, _job_id: str):
|
|
assert job_name == "process_agent_run"
|
|
assert run_id == created_run.id
|
|
assert _job_id == f"run:{created_run.id}"
|
|
assert db.committed is True
|
|
|
|
async def fake_get_arq_pool():
|
|
return Queue()
|
|
|
|
monkeypatch.setattr(agent_run_service.agent_manager, "get_agent", lambda backend_id: object())
|
|
monkeypatch.setattr(agent_run_service, "AgentRepository", AgentRepo)
|
|
monkeypatch.setattr(agent_run_service, "ConversationRepository", ConvRepo)
|
|
monkeypatch.setattr(agent_run_service, "AgentRunRepository", RunRepo)
|
|
monkeypatch.setattr(agent_run_service, "get_arq_pool", fake_get_arq_pool)
|
|
|
|
result = await agent_run_service.create_agent_run_view(
|
|
query=None,
|
|
agent_id="default",
|
|
thread_id="thread-1",
|
|
meta={"request_id": "resume-req"},
|
|
image_content=None,
|
|
current_uid="user-1",
|
|
db=db,
|
|
resume={"language": "python"},
|
|
parent_run_id="parent-run",
|
|
resume_request_id="resume-req",
|
|
)
|
|
|
|
assert result["run_id"] == created_run.id
|
|
assert db.added[0].message_type == "resume"
|
|
assert db.added[0].extra_metadata["source"] == "ask_user_question_resume"
|