wehub-resource-sync
277 lines
10 KiB
Python
277 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import uuid
|
|
from typing import Any
|
|
|
|
import asyncpg
|
|
import httpx
|
|
import pytest
|
|
|
|
pytestmark = [pytest.mark.asyncio, pytest.mark.e2e, pytest.mark.slow]
|
|
|
|
POLL_INTERVAL_SECONDS = float(os.getenv("E2E_RUN_POLL_INTERVAL_SECONDS", "2"))
|
|
RUN_TIMEOUT_SECONDS = int(os.getenv("E2E_RUN_TIMEOUT_SECONDS", "240"))
|
|
EXPECTED_OUTPUT = "ASYNC_AGENT_E2E_OK"
|
|
|
|
|
|
def _postgres_dsn() -> str:
|
|
return os.getenv("POSTGRES_URL", "postgresql+asyncpg://postgres:postgres@postgres:5432/yuxi").replace(
|
|
"+asyncpg", ""
|
|
)
|
|
|
|
|
|
async def _create_agent(client: httpx.AsyncClient, headers: dict[str, str], uid: str) -> str:
|
|
default_response = await client.get("/api/agent/default", headers=headers)
|
|
assert default_response.status_code == 200, default_response.text
|
|
default_context = ((default_response.json().get("agent") or {}).get("config_json") or {}).get("context") or {}
|
|
|
|
slug = f"e2e-async-agent-{uuid.uuid4().hex[:8]}"
|
|
context: dict[str, Any] = {
|
|
"system_prompt": f"你是端到端测试专用智能体。不要调用任何工具，只输出 {EXPECTED_OUTPUT}。",
|
|
"tools": [],
|
|
"knowledges": [],
|
|
"mcps": [],
|
|
"skills": [],
|
|
"subagents": [],
|
|
}
|
|
if default_context.get("model"):
|
|
context["model"] = default_context["model"]
|
|
|
|
response = await client.post(
|
|
"/api/agent",
|
|
json={
|
|
"name": f"E2E 异步 Agent {slug[-8:]}",
|
|
"slug": slug,
|
|
"backend_id": "ChatbotAgent",
|
|
"description": "真实异步 Agent E2E 临时智能体",
|
|
"config_json": {"context": context},
|
|
"share_config": {"access_level": "user", "department_ids": [], "user_uids": [uid]},
|
|
},
|
|
headers=headers,
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert (response.json().get("agent") or {}).get("slug") == slug
|
|
return slug
|
|
|
|
|
|
async def _delete_agent(client: httpx.AsyncClient, headers: dict[str, str], slug: str) -> None:
|
|
response = await client.delete(f"/api/agent/{slug}", headers=headers)
|
|
assert response.status_code in {200, 404}, response.text
|
|
|
|
|
|
async def _cancel_run(client: httpx.AsyncClient, headers: dict[str, str], run_id: str | None) -> None:
|
|
if not run_id:
|
|
return
|
|
response = await client.post(f"/api/agent/runs/{run_id}/cancel", headers=headers)
|
|
assert response.status_code < 500, response.text
|
|
|
|
|
|
async def _create_thread(client: httpx.AsyncClient, headers: dict[str, str], agent_slug: str) -> str:
|
|
response = await client.post(
|
|
"/api/chat/thread",
|
|
json={
|
|
"agent_id": agent_slug,
|
|
"title": f"agent-async-e2e-{uuid.uuid4().hex[:8]}",
|
|
"metadata": {"test": "agent-async-e2e"},
|
|
},
|
|
headers=headers,
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
payload = response.json()
|
|
thread_id = payload.get("thread_id") or payload.get("id")
|
|
assert thread_id, payload
|
|
return str(thread_id)
|
|
|
|
|
|
async def _create_run(
|
|
client: httpx.AsyncClient,
|
|
headers: dict[str, str],
|
|
*,
|
|
agent_slug: str,
|
|
thread_id: str,
|
|
) -> tuple[str, str]:
|
|
request_id = f"agent-async-e2e-{uuid.uuid4()}"
|
|
response = await client.post(
|
|
"/api/agent/runs",
|
|
json={
|
|
"query": f"请只回复 {EXPECTED_OUTPUT}，不要添加任何解释。",
|
|
"agent_slug": agent_slug,
|
|
"thread_id": thread_id,
|
|
"meta": {"request_id": request_id},
|
|
},
|
|
headers=headers,
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
run_id = response.json().get("run_id")
|
|
assert run_id, response.text
|
|
assert response.json().get("stream_url") == f"/api/agent/runs/{run_id}/events"
|
|
assert response.json().get("request_id") == request_id
|
|
return str(run_id), request_id
|
|
|
|
|
|
async def _iter_sse(client: httpx.AsyncClient, headers: dict[str, str], run_id: str):
|
|
async with client.stream("GET", f"/api/agent/runs/{run_id}/events?verbose=false", headers=headers) as response:
|
|
assert response.status_code == 200, response.text
|
|
event = "message"
|
|
data_lines: list[str] = []
|
|
async for line in response.aiter_lines():
|
|
if not line:
|
|
if data_lines:
|
|
yield event, json.loads("\n".join(data_lines))
|
|
event = "message"
|
|
data_lines = []
|
|
continue
|
|
if line.startswith(":"):
|
|
continue
|
|
if line.startswith("event:"):
|
|
event = line[len("event:") :].strip() or "message"
|
|
elif line.startswith("data:"):
|
|
data_lines.append(line[len("data:") :].strip())
|
|
|
|
|
|
async def _consume_events(client: httpx.AsyncClient, headers: dict[str, str], run_id: str) -> dict[str, int]:
|
|
event_counts: dict[str, int] = {}
|
|
|
|
async def consume() -> None:
|
|
async for event, payload in _iter_sse(client, headers, run_id):
|
|
event_counts[event] = event_counts.get(event, 0) + 1
|
|
if event == "end" or payload.get("status") in {"completed", "failed", "cancelled", "interrupted"}:
|
|
return
|
|
|
|
await asyncio.wait_for(consume(), timeout=RUN_TIMEOUT_SECONDS)
|
|
return event_counts
|
|
|
|
|
|
async def _wait_for_run(client: httpx.AsyncClient, headers: dict[str, str], run_id: str) -> dict:
|
|
deadline = asyncio.get_running_loop().time() + RUN_TIMEOUT_SECONDS
|
|
last_payload: dict | None = None
|
|
|
|
while asyncio.get_running_loop().time() < deadline:
|
|
response = await client.get(f"/api/agent/runs/{run_id}", headers=headers)
|
|
assert response.status_code == 200, response.text
|
|
|
|
last_payload = response.json().get("run") or {}
|
|
status = str(last_payload.get("status") or "")
|
|
if status in {"completed", "failed", "cancelled", "interrupted"}:
|
|
return last_payload
|
|
|
|
await asyncio.sleep(POLL_INTERVAL_SECONDS)
|
|
|
|
pytest.fail("Run timed out: " + json.dumps(last_payload or {}, ensure_ascii=False))
|
|
|
|
|
|
async def _assert_run_persisted(
|
|
*,
|
|
run_id: str,
|
|
request_id: str,
|
|
thread_id: str,
|
|
agent_slug: str,
|
|
uid: str,
|
|
) -> None:
|
|
conn = await asyncpg.connect(_postgres_dsn())
|
|
try:
|
|
row = await conn.fetchrow(
|
|
"""
|
|
SELECT
|
|
ar.id,
|
|
ar.request_id,
|
|
ar.status,
|
|
ar.run_type,
|
|
ar.agent_slug,
|
|
ar.uid,
|
|
ar.conversation_thread_id,
|
|
ar.conversation_id,
|
|
ar.input_message_id,
|
|
ar.output_message_id,
|
|
input_msg.role AS input_role,
|
|
input_msg.request_id AS input_request_id,
|
|
output_msg.role AS output_role,
|
|
output_msg.run_id AS output_run_id,
|
|
output_msg.content AS output_content,
|
|
conv.thread_id AS persisted_thread_id
|
|
FROM agent_runs ar
|
|
JOIN conversations conv ON conv.id = ar.conversation_id
|
|
LEFT JOIN messages input_msg ON input_msg.id = ar.input_message_id
|
|
LEFT JOIN messages output_msg ON output_msg.id = ar.output_message_id
|
|
WHERE ar.id = $1
|
|
""",
|
|
run_id,
|
|
)
|
|
assert row, f"agent_runs row missing for {run_id}"
|
|
assert row["request_id"] == request_id
|
|
assert row["status"] == "completed"
|
|
assert row["run_type"] == "chat"
|
|
assert row["agent_slug"] == agent_slug
|
|
assert row["uid"] == uid
|
|
assert row["conversation_thread_id"] == thread_id
|
|
assert row["conversation_id"] is not None
|
|
assert row["input_message_id"] is not None
|
|
assert row["output_message_id"] is not None
|
|
assert row["input_role"] == "user"
|
|
assert row["input_request_id"] == request_id
|
|
assert row["output_role"] == "assistant"
|
|
assert row["output_run_id"] == run_id
|
|
assert EXPECTED_OUTPUT in row["output_content"]
|
|
assert row["persisted_thread_id"] == thread_id
|
|
finally:
|
|
await conn.close()
|
|
|
|
|
|
async def test_async_agent_run_stream_result_and_persistence(
|
|
e2e_client: httpx.AsyncClient,
|
|
e2e_headers: dict[str, str],
|
|
e2e_agent_context: dict[str, str],
|
|
):
|
|
uid = e2e_agent_context["uid"]
|
|
agent_slug = await _create_agent(e2e_client, e2e_headers, uid)
|
|
run_id: str | None = None
|
|
run_completed = False
|
|
|
|
try:
|
|
thread_id = await _create_thread(e2e_client, e2e_headers, agent_slug)
|
|
run_id, request_id = await _create_run(
|
|
e2e_client,
|
|
e2e_headers,
|
|
agent_slug=agent_slug,
|
|
thread_id=thread_id,
|
|
)
|
|
|
|
event_counts = await _consume_events(e2e_client, e2e_headers, run_id)
|
|
assert event_counts.get("messages", 0) > 0, event_counts
|
|
assert event_counts.get("end", 0) == 1, event_counts
|
|
|
|
run_payload = await _wait_for_run(e2e_client, e2e_headers, run_id)
|
|
assert run_payload.get("status") == "completed", run_payload
|
|
assert run_payload.get("request_id") == request_id
|
|
|
|
result_response = await e2e_client.get(f"/api/agent/runs/{run_id}/result", headers=e2e_headers)
|
|
assert result_response.status_code == 200, result_response.text
|
|
result_payload = result_response.json()
|
|
assert result_payload.get("status") == "completed", result_payload
|
|
assert result_payload.get("agent_run_id") == run_id
|
|
assert result_payload.get("thread_id") == thread_id
|
|
assert result_payload.get("request_id") == request_id
|
|
assert EXPECTED_OUTPUT in str(result_payload.get("output") or ""), result_payload
|
|
|
|
history_response = await e2e_client.get(f"/api/chat/thread/{thread_id}/history", headers=e2e_headers)
|
|
assert history_response.status_code == 200, history_response.text
|
|
history_text = json.dumps(history_response.json(), ensure_ascii=False)
|
|
assert request_id in history_text, history_text
|
|
assert EXPECTED_OUTPUT in history_text, history_text
|
|
|
|
await _assert_run_persisted(
|
|
run_id=run_id,
|
|
request_id=request_id,
|
|
thread_id=thread_id,
|
|
agent_slug=agent_slug,
|
|
uid=uid,
|
|
)
|
|
run_completed = True
|
|
finally:
|
|
if not run_completed:
|
|
await _cancel_run(e2e_client, e2e_headers, run_id)
|
|
await _delete_agent(e2e_client, e2e_headers, agent_slug)
|