60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from yuxi.channels.services.context import ChatAbortEntry, ChatRunBuffer
|
||
|
|
from yuxi.channels.router import MessageRouter
|
||
|
|
|
||
|
|
|
||
|
|
class TestMessageRouterRuntime:
|
||
|
|
def test_chat_abort_controllers_initialized(self):
|
||
|
|
router = MessageRouter()
|
||
|
|
assert router.chat_abort_controllers == {}
|
||
|
|
assert router.chat_run_buffers == {}
|
||
|
|
|
||
|
|
def test_abort_chat_returns_false_for_unknown_run(self):
|
||
|
|
router = MessageRouter()
|
||
|
|
assert router.abort_chat("nonexistent") is False
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_abort_chat_cancels_running_task(self):
|
||
|
|
router = MessageRouter()
|
||
|
|
run_id = "test:msg-1"
|
||
|
|
|
||
|
|
async def slow_task():
|
||
|
|
await asyncio.sleep(10)
|
||
|
|
|
||
|
|
loop = asyncio.get_running_loop()
|
||
|
|
task = loop.create_task(slow_task())
|
||
|
|
router.chat_abort_controllers[run_id] = ChatAbortEntry(task=task)
|
||
|
|
|
||
|
|
assert router.abort_chat(run_id) is True
|
||
|
|
await asyncio.sleep(0)
|
||
|
|
assert task.cancelled()
|
||
|
|
|
||
|
|
def test_chat_run_buffer_registration(self):
|
||
|
|
router = MessageRouter()
|
||
|
|
buffer = ChatRunBuffer(run_id="test:msg-2")
|
||
|
|
router.chat_run_buffers["test:msg-2"] = buffer
|
||
|
|
|
||
|
|
buffer.append_chunk("hello")
|
||
|
|
assert router.chat_run_buffers["test:msg-2"].get_full_text() == "hello"
|
||
|
|
|
||
|
|
def test_chat_abort_entry_cleanup_on_abort(self):
|
||
|
|
router = MessageRouter()
|
||
|
|
run_id = "test:msg-3"
|
||
|
|
|
||
|
|
async def noop():
|
||
|
|
pass
|
||
|
|
|
||
|
|
loop = asyncio.new_event_loop()
|
||
|
|
task = loop.create_task(noop())
|
||
|
|
router.chat_abort_controllers[run_id] = ChatAbortEntry(task=task)
|
||
|
|
|
||
|
|
router.abort_chat(run_id)
|
||
|
|
|
||
|
|
router.chat_abort_controllers.pop(run_id, None)
|
||
|
|
assert run_id not in router.chat_abort_controllers
|
||
|
|
loop.close()
|