diff --git a/pocketstation/__init__.py b/pocketstation/__init__.py index a43259a..fd50d86 100644 --- a/pocketstation/__init__.py +++ b/pocketstation/__init__.py @@ -1,11 +1,15 @@ """PocketStation Python SDK. Phase 5.""" -from .client import PocketStation, PocketStationSession -from .types import IceServer, PocketStationError, RoomCredentials +from .client import PocketStation as PocketStationClient, PocketStationSession +from .station import PocketStation +from .types import AudioFrame, AudioMode, IceServer, PocketStationError, RoomCredentials __version__ = "0.1.0" __all__ = [ "PocketStation", + "PocketStationClient", "PocketStationSession", + "AudioFrame", + "AudioMode", "RoomCredentials", "IceServer", "PocketStationError", diff --git a/pocketstation/station.py b/pocketstation/station.py new file mode 100644 index 0000000..7b812f3 --- /dev/null +++ b/pocketstation/station.py @@ -0,0 +1,131 @@ +"""High-level PocketStation session API (spec §12.1).""" +from __future__ import annotations + +import base64 +import json +from typing import AsyncIterator, Callable, Optional + +import httpx +import websockets + +from .types import AudioFrame, AudioMode, IceServer, RoomCredentials + +_AUDIO_FRAME_MSG_TYPE = "AUDIO_FRAME" +_SUBSCRIBE_MSG_TYPE = "SUBSCRIBE" +_ROOM_STATE_MSG_TYPE = "ROOM_STATE" + + +class PocketStation: + """Voice agent / broadcast session manager (spec §12.1). + + Usage:: + + station = PocketStation(room_id="abc123", mode=AudioMode.VOICE_AGENT) + async for frame in station.listen(): + ... + """ + + def __init__( + self, + *, + room_id: Optional[str] = None, + mode: AudioMode = AudioMode.VOICE_AGENT, + api_url: str = "http://localhost:8090", + relay_url: str = "ws://localhost:8080/v1/signal", + opus_frame_duration_ms: int = 20, + on_listener_count: Optional[Callable[[int], None]] = None, + on_packet_loss: Optional[Callable[[], None]] = None, + ) -> None: + self.room_id = room_id + self.mode = mode + self.api_url = api_url + self.relay_url = relay_url + self.frame_duration_ms = opus_frame_duration_ms + self.on_listener_count = on_listener_count + self.on_packet_loss = on_packet_loss + self._ws: Optional[websockets.WebSocketClientProtocol] = None + self._credentials: Optional[RoomCredentials] = None + + async def _ensure_room(self) -> RoomCredentials: + """Create or reuse a room on the API server.""" + if self._credentials is not None: + return self._credentials + async with httpx.AsyncClient() as client: + r = await client.post(f"{self.api_url.rstrip('/')}/v1/rooms", json={}) + r.raise_for_status() + data = r.json() + creds = RoomCredentials( + room_id=data["room_id"], + source_token=data.get("source_token", ""), + listener_token=data.get("listener_token", ""), + qr_url=data.get("qr_url", ""), + ice_servers=[ + IceServer( + urls=s.get("urls", []), + username=s.get("username"), + credential=s.get("credential"), + ) + for s in data.get("ice_servers", []) + ], + ) + if not self.room_id: + self.room_id = creds.room_id + self._credentials = creds + return creds + + async def listen(self) -> AsyncIterator[AudioFrame]: + """Subscribe to audio frames from the relay. + + Connects via WebSocket SUBSCRIBE message and yields AudioFrame objects. + Falls back gracefully — yields nothing — if the relay is unreachable. + """ + creds = await self._ensure_room() + subscribe_msg = json.dumps({ + "type": _SUBSCRIBE_MSG_TYPE, + "room_id": self.room_id, + "token": creds.listener_token, + }) + try: + async with websockets.connect(self.relay_url) as ws: + self._ws = ws + await ws.send(subscribe_msg) + seq = 0 + async for raw in ws: + if isinstance(raw, bytes): + yield AudioFrame( + pcm=raw, + sequence=seq, + duration_ms=self.frame_duration_ms, + ) + seq += 1 + elif isinstance(raw, str): + msg = json.loads(raw) + if msg.get("type") == _ROOM_STATE_MSG_TYPE and self.on_listener_count: + self.on_listener_count(msg.get("listener_count", 0)) + except Exception: + # Relay unreachable or connection closed — yield nothing. + return + finally: + self._ws = None + + async def broadcast(self, audio: bytes) -> None: + """Send raw PCM bytes (f32-LE 48 kHz mono) to the relay. + + Used in voice agent response path: TTS output -> broadcast back. + Silently no-ops when not connected. + """ + if self._ws is None: + return + try: + await self._ws.send(json.dumps({ + "type": _AUDIO_FRAME_MSG_TYPE, + "pcm_b64": base64.b64encode(audio).decode(), + })) + except Exception: + pass + + async def stop(self) -> None: + """Close the WebSocket connection.""" + if self._ws is not None: + await self._ws.close() + self._ws = None diff --git a/pocketstation/types.py b/pocketstation/types.py index 17d5549..34d7000 100644 --- a/pocketstation/types.py +++ b/pocketstation/types.py @@ -1,9 +1,39 @@ """PocketStation SDK type definitions. Phase 5.""" from __future__ import annotations +import dataclasses +import enum +import struct from dataclasses import dataclass, field from typing import Optional +class AudioMode(enum.Enum): + """Audio session mode (spec §12.1).""" + VOICE = "voice" + VOICE_AGENT = "voice_agent" + MUSIC = "music" + BROADCAST = "broadcast" + + +@dataclasses.dataclass +class AudioFrame: + """A single audio frame received from the relay. + + pcm: raw PCM bytes, 48 kHz mono f32-LE. + """ + pcm: bytes + sample_rate: int = 48000 + channels: int = 1 + duration_ms: int = 20 + sequence: int = 0 + + @property + def samples(self) -> list[float]: + """Decode f32-LE PCM bytes to float samples.""" + n = len(self.pcm) // 4 + return list(struct.unpack(f"<{n}f", self.pcm[:n * 4])) + + @dataclass class IceServer: """ICE server configuration (ADR-023 embedded TURN).""" @@ -19,6 +49,7 @@ class RoomCredentials: source_token: str listener_token: str ice_servers: list[IceServer] = field(default_factory=list) + qr_url: str = "" @classmethod def from_dict(cls, data: dict) -> "RoomCredentials": @@ -35,6 +66,7 @@ def from_dict(cls, data: dict) -> "RoomCredentials": source_token=data["source_token"], listener_token=data["listener_token"], ice_servers=ice_servers, + qr_url=data.get("qr_url", ""), ) diff --git a/tests/test_client.py b/tests/test_client.py index 3bf2de9..27f0e11 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,9 +1,11 @@ """Unit tests for pocketstation.client. Phase 5.""" +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + import pytest -import httpx -import respx -from pocketstation import PocketStation, PocketStationError, RoomCredentials +from pocketstation import PocketStationClient as PocketStation, PocketStationError, RoomCredentials VALID_ROOM_RESPONSE = { @@ -12,69 +14,95 @@ "listener_token": "lst-py", } +ROOM_RESPONSE_WITH_ICE = { + **VALID_ROOM_RESPONSE, + "ice_servers": [ + {"urls": ["stun:relay.example.com:3478"]}, + {"urls": ["turn:relay.example.com:3478"], "username": "u", "credential": "p"}, + ], +} + + +def _mock_http_client(json_response: dict, status_code: int = 200): + """Return a context-manager mock for httpx.AsyncClient.""" + mock_response = MagicMock() + mock_response.is_success = status_code < 400 + mock_response.status_code = status_code + mock_response.text = "" + mock_response.json = MagicMock(return_value=json_response) + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + return mock_client + + +def _mock_http_client_error(status_code: int = 500, text: str = "internal error"): + """Return a mock httpx.AsyncClient whose response indicates an HTTP error.""" + mock_response = MagicMock() + mock_response.is_success = False + mock_response.status_code = status_code + mock_response.text = text + mock_response.json = MagicMock(return_value={}) + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + return mock_client + -@respx.mock @pytest.mark.asyncio async def test_given_valid_api_server_when_connect_then_session_has_credentials(): # Given - respx.post("http://api.example.com/v1/rooms").mock( - return_value=httpx.Response(201, json=VALID_ROOM_RESPONSE) - ) + mock_client = _mock_http_client(VALID_ROOM_RESPONSE, status_code=201) # When - async with PocketStation.connect( - api_url="http://api.example.com", - relay_url="ws://relay.example.com", - ) as session: - # Then - assert session.room_id == "room-py-001" - assert session.source_token == "src-py" - assert session.listener_token == "lst-py" - assert session.ice_servers == [] + with patch("pocketstation.client.httpx.AsyncClient", return_value=mock_client): + async with PocketStation.connect( + api_url="http://api.example.com", + relay_url="ws://relay.example.com", + ) as session: + # Then + assert session.room_id == "room-py-001" + assert session.source_token == "src-py" + assert session.listener_token == "lst-py" + assert session.ice_servers == [] -@respx.mock @pytest.mark.asyncio async def test_given_api_server_with_turn_when_connect_then_ice_servers_forwarded(): # Given - respx.post("http://api.example.com/v1/rooms").mock( - return_value=httpx.Response(201, json={ - **VALID_ROOM_RESPONSE, - "ice_servers": [ - {"urls": ["stun:relay.example.com:3478"]}, - {"urls": ["turn:relay.example.com:3478"], "username": "u", "credential": "p"}, - ], - }) - ) + mock_client = _mock_http_client(ROOM_RESPONSE_WITH_ICE, status_code=201) # When - async with PocketStation.connect( - api_url="http://api.example.com", - relay_url="ws://relay.example.com", - ) as session: - # Then - assert len(session.ice_servers) == 2 - assert session.ice_servers[1].username == "u" + with patch("pocketstation.client.httpx.AsyncClient", return_value=mock_client): + async with PocketStation.connect( + api_url="http://api.example.com", + relay_url="ws://relay.example.com", + ) as session: + # Then + assert len(session.ice_servers) == 2 + assert session.ice_servers[1].username == "u" -@respx.mock @pytest.mark.asyncio async def test_given_api_server_500_when_connect_then_raises(): # Given - respx.post("http://api.example.com/v1/rooms").mock( - return_value=httpx.Response(500, text="internal error") - ) + mock_client = _mock_http_client_error(status_code=500, text="internal error") # When / Then - with pytest.raises(PocketStationError) as exc_info: - async with PocketStation.connect( - api_url="http://api.example.com", - relay_url="ws://relay.example.com", - ): - pass + with patch("pocketstation.client.httpx.AsyncClient", return_value=mock_client): + with pytest.raises(PocketStationError) as exc_info: + async with PocketStation.connect( + api_url="http://api.example.com", + relay_url="ws://relay.example.com", + ): + pass assert exc_info.value.code == "http_error" @pytest.mark.asyncio async def test_given_preexisting_credentials_when_connect_then_no_http_call(): - # Given — no network mock; any http call would error + # Given — no patch; any real HTTP call would fail with PermissionError in sandbox creds = RoomCredentials.from_dict(VALID_ROOM_RESPONSE) # When async with PocketStation.connect( diff --git a/tests/test_station.py b/tests/test_station.py new file mode 100644 index 0000000..284e603 --- /dev/null +++ b/tests/test_station.py @@ -0,0 +1,144 @@ +"""Unit tests for pocketstation.station — spec §12.1 voice agent pattern.""" +from __future__ import annotations + +import struct +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from pocketstation import AudioFrame, AudioMode, PocketStation +from pocketstation.types import RoomCredentials + + +VALID_ROOM_RESPONSE = { + "room_id": "room-station-001", + "source_token": "src-station", + "listener_token": "lst-station", +} + + +def _make_pcm(n_samples: int = 4) -> bytes: + """Build n_samples of f32-LE PCM silence.""" + return struct.pack(f"<{n_samples}f", *([0.0] * n_samples)) + + +def _mock_http_client(json_response: dict, status_code: int = 200): + """Return a context-manager mock for httpx.AsyncClient that yields a mock response.""" + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json = MagicMock(return_value=json_response) + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + return mock_client + + +# --------------------------------------------------------------------------- +# AudioMode / construction +# --------------------------------------------------------------------------- + + +def test_given_voice_agent_mode_when_created_then_fields_set(): + # Given / When + station = PocketStation(room_id="abc123", mode=AudioMode.VOICE_AGENT) + # Then + assert station.room_id == "abc123" + assert station.mode is AudioMode.VOICE_AGENT + assert station.frame_duration_ms == 20 + assert station._ws is None + assert station._credentials is None + + +# --------------------------------------------------------------------------- +# AudioFrame.samples property +# --------------------------------------------------------------------------- + + +def test_given_audio_frame_when_samples_then_correct_count(): + # Given + n = 8 + pcm = _make_pcm(n) + frame = AudioFrame(pcm=pcm) + # When + samples = frame.samples + # Then + assert len(samples) == n + assert all(s == 0.0 for s in samples) + + +def test_given_audio_frame_with_known_values_when_samples_then_decoded_correctly(): + # Given + values = [0.5, -0.5, 1.0, -1.0] + pcm = struct.pack("<4f", *values) + frame = AudioFrame(pcm=pcm) + # When + samples = frame.samples + # Then + assert len(samples) == 4 + for actual, expected in zip(samples, values): + assert abs(actual - expected) < 1e-6 + + +# --------------------------------------------------------------------------- +# _ensure_room — credential reuse +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_given_room_id_when_ensure_room_then_reuses_credentials(): + # Given — first call creates room via mocked HTTP client + mock_client = _mock_http_client(VALID_ROOM_RESPONSE) + + station = PocketStation(room_id="existing-room") + + with patch("pocketstation.station.httpx.AsyncClient", return_value=mock_client): + # When — call _ensure_room twice + creds1 = await station._ensure_room() + creds2 = await station._ensure_room() + + # Then — same object returned; HTTP was called only once + assert creds1 is creds2 + assert creds1.room_id == "room-station-001" + mock_client.post.assert_called_once() + + +@pytest.mark.asyncio +async def test_given_no_room_id_when_ensure_room_then_room_id_assigned(): + # Given + mock_client = _mock_http_client(VALID_ROOM_RESPONSE) + + station = PocketStation() # no room_id + + with patch("pocketstation.station.httpx.AsyncClient", return_value=mock_client): + # When + creds = await station._ensure_room() + + # Then + assert station.room_id == "room-station-001" + assert creds.room_id == "room-station-001" + + +# --------------------------------------------------------------------------- +# listen() — websocket unavailable +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_given_websocket_unavailable_when_listen_then_yields_nothing(): + # Given — HTTP room creation succeeds, WebSocket connect raises immediately + mock_client = _mock_http_client(VALID_ROOM_RESPONSE) + + station = PocketStation(room_id="abc123", relay_url="ws://127.0.0.1:1") + + with patch("pocketstation.station.httpx.AsyncClient", return_value=mock_client), \ + patch("pocketstation.station.websockets.connect", side_effect=OSError("unreachable")): + # When + frames = [] + async for frame in station.listen(): + frames.append(frame) + + # Then — no frames, no exception raised + assert frames == []