forked from agentclientprotocol/python-sdk
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_core.py
More file actions
55 lines (41 loc) · 1.54 KB
/
Copy pathtest_core.py
File metadata and controls
55 lines (41 loc) · 1.54 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
from __future__ import annotations
import asyncio
import contextlib
from typing import Any
import pytest
from acp.core import run_agent
@pytest.mark.asyncio
async def test_run_agent_closes_connection_when_cancelled(server, agent) -> None:
sender_created = asyncio.Event()
sender_closed = asyncio.Event()
dispatcher_started = asyncio.Event()
dispatcher_stopped = asyncio.Event()
class TrackingSender:
def __init__(self, writer: asyncio.StreamWriter, supervisor: Any) -> None:
sender_created.set()
async def send(self, payload: dict[str, Any]) -> None:
msg = "test does not send messages"
raise AssertionError(msg)
async def close(self) -> None:
sender_closed.set()
class TrackingDispatcher:
def start(self) -> None:
dispatcher_started.set()
async def stop(self) -> None:
dispatcher_stopped.set()
task = asyncio.create_task(
run_agent(
agent,
server.server_writer,
server.server_reader,
sender_factory=TrackingSender,
dispatcher_factory=lambda *args: TrackingDispatcher(),
)
)
await asyncio.wait_for(sender_created.wait(), timeout=1)
await asyncio.wait_for(dispatcher_started.wait(), timeout=1)
task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
await asyncio.wait_for(dispatcher_stopped.wait(), timeout=1)
await asyncio.wait_for(sender_closed.wait(), timeout=1)