43 lines
1.5 KiB
Python
43 lines
1.5 KiB
Python
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()
|