fix agent idle stream heartbeat visibility
This commit is contained in:
@@ -0,0 +1,42 @@
|
||||
import asyncio
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
_TEST_DATA_DIR = None
|
||||
if "SOVEREIGNAI_DATA_DIR" not in os.environ:
|
||||
_TEST_DATA_DIR = tempfile.TemporaryDirectory(prefix="sovereignai-stream-tests-")
|
||||
os.environ["SOVEREIGNAI_DATA_DIR"] = _TEST_DATA_DIR.name
|
||||
|
||||
from app import main
|
||||
|
||||
|
||||
class AgentStreamTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_idle_model_wait_sends_heartbeat_then_final_result(self) -> None:
|
||||
async def slow_agent(request, report_progress=None, *, user_id):
|
||||
await report_progress("بدأ تحليل المهمة.")
|
||||
await asyncio.sleep(0.05)
|
||||
return {"task": request.task, "result": "اكتمل التحليل."}
|
||||
|
||||
with (
|
||||
patch.object(main, "AGENT_STREAM_HEARTBEAT_SECONDS", 0.01),
|
||||
patch.object(main, "_execute_agent", side_effect=slow_agent),
|
||||
):
|
||||
response = await main.run_agent_stream(
|
||||
main.AgentRequest(task="اختبار انتظار النموذج"), user_id="test-user"
|
||||
)
|
||||
chunks = [chunk async for chunk in response.body_iterator]
|
||||
|
||||
body = b"".join(
|
||||
chunk if isinstance(chunk, bytes) else chunk.encode("utf-8")
|
||||
for chunk in chunks
|
||||
).decode("utf-8")
|
||||
self.assertIn("event: progress", body)
|
||||
self.assertIn("event: heartbeat\ndata: {}", body)
|
||||
self.assertIn("event: done", body)
|
||||
self.assertIn("اكتمل التحليل", body)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user