Assistant/client/tests/test_session.py

import asyncio
import json

import pytest

from client.session import Session


class FakeWebSocket:
    """Minimal stand-in for a `websockets` client connection."""
    def __init__(self, inbox: list):
        # inbox is a list of either str (text frame) or bytes (binary frame).
        self._inbox = list(inbox)
        self.sent: list = []

    async def send(self, frame):
        self.sent.append(frame)

    async def recv(self):
        if not self._inbox:
            await asyncio.sleep(3600)  # block forever; tests cancel
        return self._inbox.pop(0)

    async def close(self):
        pass


class _FakeCallbacks:
    """Implements client.session.Callbacks via duck-typing."""
    def __init__(self):
        self.started: list = []
        self.audio: list = []
        self.assistant_done = 0
        self.session_ended: list = []

    def on_session_started(self, conversation_id: str) -> None:
        self.started.append(conversation_id)

    def on_assistant_audio(self, pcm16: bytes) -> None:
        self.audio.append(pcm16)

    def on_assistant_done(self) -> None:
        self.assistant_done += 1

    def on_session_ended(self, reason: str) -> None:
        self.session_ended.append(reason)


async def _run_pump_briefly(session, ws, settle: float = 0.05):
    task = asyncio.create_task(session._pump(ws))
    await asyncio.sleep(settle)
    task.cancel()
    try:
        await task
    except (asyncio.CancelledError, BaseException):
        pass


async def test_hello_is_sent_on_connect_and_hello_ack_logged():
    ws = FakeWebSocket(inbox=[
        json.dumps({"type": "hello_ack", "config": {"voice": "alloy"}}),
    ])
    session = Session(
        backend_ws="wss://x/device",
        device_token="T",
        device_id="00000000-0000-0000-0000-000000000001",
        callbacks=_FakeCallbacks(),
    )
    await _run_pump_briefly(session, ws)

    first = json.loads(ws.sent[0])
    assert first["type"] == "hello"
    assert first["device_id"] == "00000000-0000-0000-0000-000000000001"
    assert first["client_version"]


async def test_session_started_calls_callback_with_conversation_id():
    cb = _FakeCallbacks()
    ws = FakeWebSocket(inbox=[
        json.dumps({"type": "hello_ack", "config": {}}),
        json.dumps({
            "type": "session_started",
            "conversation_id": "cccccccc-cccc-cccc-cccc-cccccccccccc",
        }),
    ])
    session = Session("wss://x/device", "T", "did", cb)
    await _run_pump_briefly(session, ws)
    assert cb.started == ["cccccccc-cccc-cccc-cccc-cccccccccccc"]


async def test_binary_frame_goes_to_audio_callback():
    cb = _FakeCallbacks()
    payload = bytes(range(256)) * 15  # 3840 bytes
    ws = FakeWebSocket(inbox=[
        json.dumps({"type": "hello_ack", "config": {}}),
        payload,
    ])
    session = Session("wss://x/device", "T", "did", cb)
    await _run_pump_briefly(session, ws)
    assert cb.audio == [payload]


async def test_assistant_done_and_session_ended_callbacks():
    cb = _FakeCallbacks()
    ws = FakeWebSocket(inbox=[
        json.dumps({"type": "hello_ack", "config": {}}),
        json.dumps({"type": "assistant_done"}),
        json.dumps({"type": "session_ended", "reason": "idle"}),
    ])
    session = Session("wss://x/device", "T", "did", cb)
    await _run_pump_briefly(session, ws)
    assert cb.assistant_done == 1
    assert cb.session_ended == ["idle"]


async def test_tool_call_runs_handler_and_sends_tool_result():
    cb = _FakeCallbacks()
    ws = FakeWebSocket(inbox=[
        json.dumps({"type": "hello_ack", "config": {}}),
        json.dumps({
            "type": "tool_call",
            "call_id": "c1",
            "name": "set_volume",
            "arguments": {"level": 42},
        }),
    ])
    session = Session("wss://x/device", "T", "did", cb)
    captured = {}

    def handler(args):
        captured["args"] = args
        return {"ok": True, "level": args["level"]}

    session.register_tool("set_volume", handler)
    await _run_pump_briefly(session, ws)

    assert captured["args"] == {"level": 42}
    reply = json.loads(ws.sent[-1])
    assert reply["type"] == "tool_result"
    assert reply["call_id"] == "c1"
    assert reply["ok"] is True
    assert reply["result"] == {"ok": True, "level": 42}


async def test_tool_call_with_exception_sends_ok_false():
    cb = _FakeCallbacks()
    ws = FakeWebSocket(inbox=[
        json.dumps({"type": "hello_ack", "config": {}}),
        json.dumps({
            "type": "tool_call",
            "call_id": "c2",
            "name": "broken",
            "arguments": {},
        }),
    ])
    session = Session("wss://x/device", "T", "did", cb)

    def handler(args):
        raise RuntimeError("kaboom")

    session.register_tool("broken", handler)
    await _run_pump_briefly(session, ws)

    reply = json.loads(ws.sent[-1])
    assert reply["type"] == "tool_result"
    assert reply["call_id"] == "c2"
    assert reply["ok"] is False
    assert "kaboom" in reply["error"]


async def test_unknown_tool_sends_ok_false():
    cb = _FakeCallbacks()
    ws = FakeWebSocket(inbox=[
        json.dumps({"type": "hello_ack", "config": {}}),
        json.dumps({
            "type": "tool_call",
            "call_id": "c3",
            "name": "ghost",
            "arguments": {},
        }),
    ])
    session = Session("wss://x/device", "T", "did", cb)
    await _run_pump_briefly(session, ws)
    reply = json.loads(ws.sent[-1])
    assert reply["ok"] is False
    assert "ghost" in reply["error"]