beaver test
This commit is contained in:
0Xiao0
426
test_beaver_terminal_client.py
Normal file
426
test_beaver_terminal_client.py
Normal file
@ -0,0 +1,426 @@
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
raise SystemExit(pytest.main([__file__]))
|
||||
|
||||
try:
|
||||
from custom.beaver_terminal_client import (
|
||||
BeaverTerminalClient,
|
||||
BeaverTerminalError,
|
||||
MessageIdGenerator,
|
||||
build_connect_frame,
|
||||
build_message_frame,
|
||||
)
|
||||
except ModuleNotFoundError:
|
||||
from beaver_terminal_client import (
|
||||
BeaverTerminalClient,
|
||||
BeaverTerminalError,
|
||||
MessageIdGenerator,
|
||||
build_connect_frame,
|
||||
build_message_frame,
|
||||
)
|
||||
|
||||
|
||||
def test_build_connect_frame_uses_stable_peer_id() -> None:
|
||||
frame = build_connect_frame(peer_id="device-001", device_name="desk-terminal")
|
||||
|
||||
assert frame == {
|
||||
"type": "connect",
|
||||
"peer_id": "device-001",
|
||||
"device_name": "desk-terminal",
|
||||
"capabilities": ["text"],
|
||||
}
|
||||
|
||||
|
||||
def test_build_message_frame_uses_message_id_and_text() -> None:
|
||||
frame = build_message_frame(message_id="device-001-000001", text="hello")
|
||||
|
||||
assert frame == {
|
||||
"type": "message",
|
||||
"message_id": "device-001-000001",
|
||||
"text": "hello",
|
||||
}
|
||||
|
||||
|
||||
def test_message_id_generator_uses_monotonic_peer_counter() -> None:
|
||||
generator = MessageIdGenerator(peer_id="device-001", initial_counter=7)
|
||||
|
||||
assert generator.next_id() == "device-001-000008"
|
||||
assert generator.next_id() == "device-001-000009"
|
||||
assert generator.counter == 9
|
||||
|
||||
|
||||
def test_message_id_generator_can_include_instance_id() -> None:
|
||||
generator = MessageIdGenerator(peer_id="device-001", instance_id="abc123ef")
|
||||
|
||||
assert generator.next_id() == "device-001-abc123ef-000001"
|
||||
assert generator.next_id() == "device-001-abc123ef-000002"
|
||||
|
||||
|
||||
async def test_client_connects_sends_text_and_returns_assistant_reply(
|
||||
unused_tcp_port: int,
|
||||
) -> None:
|
||||
received: list[dict[str, object]] = []
|
||||
|
||||
async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
|
||||
ws = web.WebSocketResponse()
|
||||
await ws.prepare(request)
|
||||
|
||||
async for message in ws:
|
||||
assert message.type == aiohttp.WSMsgType.TEXT
|
||||
frame = json.loads(message.data)
|
||||
received.append(frame)
|
||||
|
||||
if frame["type"] == "connect":
|
||||
await ws.send_json(
|
||||
{
|
||||
"type": "connected",
|
||||
"channel_id": "terminal-dev",
|
||||
"session_id": "terminal-dev:local:device-001",
|
||||
}
|
||||
)
|
||||
elif frame["type"] == "message":
|
||||
await ws.send_json(
|
||||
{
|
||||
"type": "ack",
|
||||
"message_id": frame["message_id"],
|
||||
"session_id": "terminal-dev:local:device-001",
|
||||
"accepted": True,
|
||||
}
|
||||
)
|
||||
await ws.send_json(
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"message_id": frame["message_id"],
|
||||
"run_id": "run-1",
|
||||
"text": "assistant reply",
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
)
|
||||
|
||||
return ws
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_get("/api/channels/terminal-dev/ws", websocket_handler)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", unused_tcp_port)
|
||||
await site.start()
|
||||
|
||||
client = BeaverTerminalClient(
|
||||
url=f"http://127.0.0.1:{unused_tcp_port}/api/channels/terminal-dev/ws",
|
||||
peer_id="device-001",
|
||||
device_name="desk-terminal",
|
||||
message_ids=MessageIdGenerator(peer_id="device-001"),
|
||||
)
|
||||
|
||||
try:
|
||||
await client.connect()
|
||||
reply = await client.send_text("hello")
|
||||
finally:
|
||||
await client.close()
|
||||
await runner.cleanup()
|
||||
|
||||
assert client.session_id == "terminal-dev:local:device-001"
|
||||
assert reply == "assistant reply"
|
||||
assert received == [
|
||||
{
|
||||
"type": "connect",
|
||||
"peer_id": "device-001",
|
||||
"device_name": "desk-terminal",
|
||||
"capabilities": ["text"],
|
||||
},
|
||||
{
|
||||
"type": "message",
|
||||
"message_id": "device-001-000001",
|
||||
"text": "hello",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
async def test_client_returns_cached_duplicate_reply(unused_tcp_port: int) -> None:
|
||||
async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
|
||||
ws = web.WebSocketResponse()
|
||||
await ws.prepare(request)
|
||||
|
||||
async for message in ws:
|
||||
frame = json.loads(message.data)
|
||||
if frame["type"] == "connect":
|
||||
await ws.send_json(
|
||||
{
|
||||
"type": "connected",
|
||||
"channel_id": "terminal-dev",
|
||||
"session_id": "terminal-dev:local:device-001",
|
||||
}
|
||||
)
|
||||
elif frame["type"] == "message":
|
||||
await ws.send_json(
|
||||
{
|
||||
"type": "ack",
|
||||
"message_id": frame["message_id"],
|
||||
"session_id": "terminal-dev:local:device-001",
|
||||
"accepted": False,
|
||||
"duplicate": True,
|
||||
"pending": False,
|
||||
"reply": "cached assistant reply",
|
||||
}
|
||||
)
|
||||
|
||||
return ws
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_get("/api/channels/terminal-dev/ws", websocket_handler)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", unused_tcp_port)
|
||||
await site.start()
|
||||
|
||||
client = BeaverTerminalClient(
|
||||
url=f"http://127.0.0.1:{unused_tcp_port}/api/channels/terminal-dev/ws",
|
||||
peer_id="device-001",
|
||||
device_name="desk-terminal",
|
||||
message_ids=MessageIdGenerator(peer_id="device-001"),
|
||||
)
|
||||
|
||||
try:
|
||||
await client.connect()
|
||||
reply = await client.send_text("hello")
|
||||
finally:
|
||||
await client.close()
|
||||
await runner.cleanup()
|
||||
|
||||
assert reply == "cached assistant reply"
|
||||
|
||||
|
||||
async def test_client_raises_on_error_frames(unused_tcp_port: int) -> None:
|
||||
async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
|
||||
ws = web.WebSocketResponse()
|
||||
await ws.prepare(request)
|
||||
|
||||
async for message in ws:
|
||||
frame = json.loads(message.data)
|
||||
if frame["type"] == "connect":
|
||||
await ws.send_json(
|
||||
{
|
||||
"type": "connected",
|
||||
"channel_id": "terminal-dev",
|
||||
"session_id": "terminal-dev:local:device-001",
|
||||
}
|
||||
)
|
||||
elif frame["type"] == "message":
|
||||
await ws.send_json({"type": "error", "error": "text is required"})
|
||||
|
||||
return ws
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_get("/api/channels/terminal-dev/ws", websocket_handler)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", unused_tcp_port)
|
||||
await site.start()
|
||||
|
||||
client = BeaverTerminalClient(
|
||||
url=f"http://127.0.0.1:{unused_tcp_port}/api/channels/terminal-dev/ws",
|
||||
peer_id="device-001",
|
||||
device_name="desk-terminal",
|
||||
message_ids=MessageIdGenerator(peer_id="device-001"),
|
||||
)
|
||||
|
||||
try:
|
||||
await client.connect()
|
||||
with pytest.raises(BeaverTerminalError, match="text is required"):
|
||||
await client.send_text("hello")
|
||||
finally:
|
||||
await client.close()
|
||||
await runner.cleanup()
|
||||
|
||||
|
||||
async def test_client_treats_assistant_finish_reason_error_as_failed_turn(
|
||||
unused_tcp_port: int,
|
||||
) -> None:
|
||||
async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
|
||||
ws = web.WebSocketResponse()
|
||||
await ws.prepare(request)
|
||||
|
||||
async for message in ws:
|
||||
frame = json.loads(message.data)
|
||||
if frame["type"] == "connect":
|
||||
await ws.send_json(
|
||||
{
|
||||
"type": "connected",
|
||||
"channel_id": "terminal-dev",
|
||||
"session_id": "terminal-dev:local:device-001",
|
||||
}
|
||||
)
|
||||
elif frame["type"] == "message":
|
||||
await ws.send_json(
|
||||
{
|
||||
"type": "ack",
|
||||
"message_id": frame["message_id"],
|
||||
"session_id": "terminal-dev:local:device-001",
|
||||
"accepted": True,
|
||||
}
|
||||
)
|
||||
await ws.send_json(
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"message_id": frame["message_id"],
|
||||
"run_id": "run-1",
|
||||
"text": "failed turn",
|
||||
"finish_reason": "error",
|
||||
}
|
||||
)
|
||||
|
||||
return ws
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_get("/api/channels/terminal-dev/ws", websocket_handler)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", unused_tcp_port)
|
||||
await site.start()
|
||||
|
||||
client = BeaverTerminalClient(
|
||||
url=f"http://127.0.0.1:{unused_tcp_port}/api/channels/terminal-dev/ws",
|
||||
peer_id="device-001",
|
||||
device_name="desk-terminal",
|
||||
message_ids=MessageIdGenerator(peer_id="device-001"),
|
||||
)
|
||||
|
||||
try:
|
||||
await client.connect()
|
||||
with pytest.raises(BeaverTerminalError, match="failed turn"):
|
||||
await client.send_text("hello")
|
||||
finally:
|
||||
await client.close()
|
||||
await runner.cleanup()
|
||||
|
||||
|
||||
async def test_client_ping_sends_ping_and_waits_for_pong(unused_tcp_port: int) -> None:
|
||||
async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
|
||||
ws = web.WebSocketResponse()
|
||||
await ws.prepare(request)
|
||||
|
||||
async for message in ws:
|
||||
frame = json.loads(message.data)
|
||||
if frame["type"] == "connect":
|
||||
await ws.send_json(
|
||||
{
|
||||
"type": "connected",
|
||||
"channel_id": "terminal-dev",
|
||||
"session_id": "terminal-dev:local:device-001",
|
||||
}
|
||||
)
|
||||
elif frame["type"] == "ping":
|
||||
await ws.send_json({"type": "pong"})
|
||||
|
||||
return ws
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_get("/api/channels/terminal-dev/ws", websocket_handler)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", unused_tcp_port)
|
||||
await site.start()
|
||||
|
||||
client = BeaverTerminalClient(
|
||||
url=f"http://127.0.0.1:{unused_tcp_port}/api/channels/terminal-dev/ws",
|
||||
peer_id="device-001",
|
||||
device_name="desk-terminal",
|
||||
message_ids=MessageIdGenerator(peer_id="device-001"),
|
||||
)
|
||||
|
||||
try:
|
||||
await client.connect()
|
||||
assert await client.ping()
|
||||
finally:
|
||||
await client.close()
|
||||
await runner.cleanup()
|
||||
|
||||
|
||||
async def test_client_reconnects_with_same_peer_id_when_socket_closes_before_send(
|
||||
unused_tcp_port: int,
|
||||
) -> None:
|
||||
connect_peer_ids: list[str] = []
|
||||
message_ids: list[str] = []
|
||||
connection_count = 0
|
||||
|
||||
async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
|
||||
nonlocal connection_count
|
||||
connection_count += 1
|
||||
current_connection = connection_count
|
||||
ws = web.WebSocketResponse()
|
||||
await ws.prepare(request)
|
||||
|
||||
async for message in ws:
|
||||
frame = json.loads(message.data)
|
||||
if frame["type"] == "connect":
|
||||
connect_peer_ids.append(frame["peer_id"])
|
||||
await ws.send_json(
|
||||
{
|
||||
"type": "connected",
|
||||
"channel_id": "terminal-dev",
|
||||
"session_id": "terminal-dev:local:device-001",
|
||||
}
|
||||
)
|
||||
elif frame["type"] == "message":
|
||||
message_ids.append(frame["message_id"])
|
||||
if current_connection == 1:
|
||||
await ws.close()
|
||||
continue
|
||||
await ws.send_json(
|
||||
{
|
||||
"type": "ack",
|
||||
"message_id": frame["message_id"],
|
||||
"session_id": "terminal-dev:local:device-001",
|
||||
"accepted": True,
|
||||
}
|
||||
)
|
||||
await ws.send_json(
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"message_id": frame["message_id"],
|
||||
"run_id": "run-2",
|
||||
"text": "reply after reconnect",
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
)
|
||||
|
||||
return ws
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_get("/api/channels/terminal-dev/ws", websocket_handler)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", unused_tcp_port)
|
||||
await site.start()
|
||||
|
||||
client = BeaverTerminalClient(
|
||||
url=f"http://127.0.0.1:{unused_tcp_port}/api/channels/terminal-dev/ws",
|
||||
peer_id="device-001",
|
||||
device_name="desk-terminal",
|
||||
message_ids=MessageIdGenerator(peer_id="device-001"),
|
||||
)
|
||||
|
||||
try:
|
||||
await client.connect()
|
||||
await asyncio.sleep(0.01)
|
||||
reply = await client.send_text("hello")
|
||||
finally:
|
||||
await client.close()
|
||||
await runner.cleanup()
|
||||
|
||||
assert reply == "reply after reconnect"
|
||||
assert connect_peer_ids == ["device-001", "device-001"]
|
||||
assert message_ids == ["device-001-000001", "device-001-000002"]
|
||||
Reference in New Issue
Block a user