632ea8c032
Co-authored-by: Antoine <33425718+anth2o@users.noreply.github.com> Co-authored-by: Bastien <bastien.baret@gmail.com> Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai> Co-authored-by: Kim-Adeline Miguel <51720070+kimadeline@users.noreply.github.com> Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai> Co-authored-by: Maxime Dolores <maxime.dolores@ext.mistral.ai> Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com> Co-authored-by: Nelson PROIA <144663685+Nelson-PROIA@users.noreply.github.com> Co-authored-by: Pierre Rossinès <pierre.rossines@mistral.ai> Co-authored-by: Quentin <quentin.torroba@mistral.ai> Co-authored-by: Robin Gullo <robin.gullo@mistral.ai> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
155 lines
3.9 KiB
Python
155 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from vibe.acp.session import AcpSessionLoop
|
|
|
|
|
|
def _make_session() -> AcpSessionLoop:
|
|
return AcpSessionLoop(
|
|
id="test-session", agent_loop=MagicMock(), command_registry=MagicMock()
|
|
)
|
|
|
|
|
|
class TestSpawn:
|
|
@pytest.mark.asyncio
|
|
async def test_spawn_creates_task(self) -> None:
|
|
session = _make_session()
|
|
ran = asyncio.Event()
|
|
|
|
async def work() -> None:
|
|
ran.set()
|
|
|
|
task = session.spawn(work())
|
|
|
|
assert task is not None
|
|
await task
|
|
assert ran.is_set()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_spawn_returns_none_after_close(self) -> None:
|
|
session = _make_session()
|
|
await session.close()
|
|
|
|
async def noop() -> None:
|
|
pass
|
|
|
|
assert session.spawn(noop()) is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_spawn_tracks_multiple_tasks(self) -> None:
|
|
session = _make_session()
|
|
gate = asyncio.Event()
|
|
|
|
async def wait_for_gate() -> None:
|
|
await gate.wait()
|
|
|
|
t1 = session.spawn(wait_for_gate())
|
|
t2 = session.spawn(wait_for_gate())
|
|
|
|
assert t1 is not None
|
|
assert t2 is not None
|
|
assert not t1.done()
|
|
assert not t2.done()
|
|
|
|
gate.set()
|
|
await asyncio.gather(t1, t2)
|
|
|
|
|
|
class TestPromptTask:
|
|
@pytest.mark.asyncio
|
|
async def test_set_prompt_task_tracks_task(self) -> None:
|
|
session = _make_session()
|
|
|
|
async def work() -> None:
|
|
pass
|
|
|
|
task = session.set_prompt_task(work())
|
|
assert session.prompt_task is task
|
|
await task
|
|
assert session.prompt_task is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_prompt_cancels_active_task(self) -> None:
|
|
session = _make_session()
|
|
|
|
async def hang() -> None:
|
|
await asyncio.Event().wait()
|
|
|
|
task = session.set_prompt_task(hang())
|
|
await session.cancel_prompt()
|
|
assert task.cancelled()
|
|
assert session.prompt_task is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_prompt_is_noop_without_task(self) -> None:
|
|
session = _make_session()
|
|
await session.cancel_prompt()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_prompt_does_not_cancel_background_tasks(self) -> None:
|
|
session = _make_session()
|
|
gate = asyncio.Event()
|
|
|
|
async def bg() -> None:
|
|
await gate.wait()
|
|
|
|
bg_task = session.spawn(bg())
|
|
assert bg_task is not None
|
|
|
|
async def hang() -> None:
|
|
await asyncio.Event().wait()
|
|
|
|
session.set_prompt_task(hang())
|
|
await session.cancel_prompt()
|
|
|
|
assert not bg_task.cancelled()
|
|
gate.set()
|
|
await bg_task
|
|
|
|
|
|
class TestClose:
|
|
@pytest.mark.asyncio
|
|
async def test_close_cancels_all_tasks(self) -> None:
|
|
session = _make_session()
|
|
|
|
async def hang() -> None:
|
|
await asyncio.Event().wait()
|
|
|
|
bg = session.spawn(hang())
|
|
prompt = session.set_prompt_task(hang())
|
|
|
|
assert bg is not None
|
|
|
|
await session.close()
|
|
|
|
assert bg.cancelled()
|
|
assert prompt.cancelled()
|
|
assert session.prompt_task is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_is_idempotent(self) -> None:
|
|
session = _make_session()
|
|
await session.close()
|
|
await session.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_waits_for_task_cleanup(self) -> None:
|
|
session = _make_session()
|
|
cleanup_ran = asyncio.Event()
|
|
|
|
async def task_with_cleanup() -> None:
|
|
try:
|
|
await asyncio.Event().wait()
|
|
except asyncio.CancelledError:
|
|
cleanup_ran.set()
|
|
raise
|
|
|
|
session.spawn(task_with_cleanup())
|
|
await asyncio.sleep(0) # let the task start
|
|
await session.close()
|
|
|
|
assert cleanup_ran.is_set()
|