mcp_screen/hermes_voice_gateway/test_websocket.py

# pyright: reportMissingImports=false, reportAttributeAccessIssue=false
"""Integration and unit tests for the ESP32 Voice WebSocket gateway endpoint."""

import asyncio
import json
import logging
import sys
import unittest
from pathlib import Path
from unittest.mock import MagicMock

# 1. Mock external tools modules before importing the API endpoint
mock_transcription = MagicMock()
mock_transcription.transcribe_audio.return_value = {"success": True, "transcript": "Hello world"}
sys.modules["tools.transcription_tools"] = mock_transcription

mock_tts = MagicMock()
mock_tts.text_to_speech_tool.return_value = json.dumps({"success": True, "file_path": __file__})
sys.modules["tools.tts_tool"] = mock_tts

# Mock aiohttp
from aiohttp import web, ClientSession

# Import gateway mixin and register function
sys.path.append(str(Path(__file__).parent))
from api_server_endpoint import HermesESP32VoiceGatewayMixin, register_esp32_voice_route

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class MockHermesServer(HermesESP32VoiceGatewayMixin):
    def __init__(self):
        self.auth_token = "valid_secret_key"
        self._background_tasks = set()

    def _check_auth(self, request):
        auth_header = request.headers.get("Authorization")
        if not auth_header or auth_header != f"Bearer {self.auth_token}":
            return web.json_response({"error": "Unauthorized"}, status=401)
        return None

    async def _run_esp32_voice_turn(self, transcript, device_id, instructions=None):
        return {"id": "resp_test123"}, f"Echo: {transcript}"

    @staticmethod
    def _convert_audio_to_esp32_wav(input_path, *, output_dir, device_id):
        out_path = Path(output_dir) / "mock_response.wav"
        import wave
        with wave.open(str(out_path), "wb") as wav:
            wav.setnchannels(1)
            wav.setsampwidth(2)
            wav.setframerate(16000)
            wav.writeframes(b"\x00" * 8000)  # Dummy PCM data
        return out_path


class TestWebSocketGateway(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.server = MockHermesServer()
        self.app = web.Application()
        register_esp32_voice_route(self.app, self.server)
        
        self.runner = web.AppRunner(self.app)
        await self.runner.setup()
        self.site = web.TCPSite(self.runner, "127.0.0.1", 0)
        await self.site.start()
        
        self.port = self.runner.addresses[0][1]
        self.base_url = f"http://127.0.0.1:{self.port}"
        self.ws_url = f"ws://127.0.0.1:{self.port}/api/esp32/voice/ws"

    async def asyncTearDown(self):
        await self.runner.cleanup()

    async def test_auth_failure_missing_token(self):
        """Verify authentication failure (missing token) returns 401."""
        async with ClientSession() as session:
            try:
                async with session.ws_connect(self.ws_url) as ws:
                    pass
                self.fail("Connection should have been rejected with 401")
            except Exception as e:
                # HTTP status code should be checked
                pass

            # Test HTTP request directly to assert 401
            async with session.get(self.ws_url) as resp:
                self.assertEqual(resp.status, 401)
                body = await resp.json()
                self.assertEqual(body["error"], "Unauthorized")

    async def test_auth_failure_bad_token(self):
        """Verify authentication failure (bad token) returns 401."""
        headers = {"Authorization": "Bearer bad_token"}
        async with ClientSession() as session:
            async with session.get(self.ws_url, headers=headers) as resp:
                self.assertEqual(resp.status, 401)

    async def test_successful_streaming_flow(self):
        """Verify a normal WebSocket start, stream, stop, and response flow."""
        headers = {
            "Authorization": "Bearer valid_secret_key",
            "X-Device-ID": "test-device"
        }
        
        async with ClientSession() as session:
            async with session.ws_connect(self.ws_url, headers=headers) as ws:
                # 1. Send start event
                await ws.send_str(json.dumps({
                    "event": "start",
                    "device_id": "test-device",
                    "sample_rate": 16000
                }))

                # Read events
                ready_msg = await ws.receive_json()
                self.assertEqual(ready_msg["event"], "ready")

                listening_msg = await ws.receive_json()
                self.assertEqual(listening_msg["event"], "listening")

                # 2. Send simulated binary PCM audio chunks (Total: 4000 bytes, > 3200 byte limit)
                await ws.send_bytes(b"\x00" * 2000)
                await ws.send_bytes(b"\x00" * 2000)

                # 3. Send stop event
                await ws.send_str(json.dumps({"event": "stop"}))

                # 4. Check transcription event
                transcript_msg = await ws.receive_json()
                self.assertEqual(transcript_msg["event"], "transcript")
                self.assertEqual(transcript_msg["text"], "Hello world")

                # 5. Check thinking event
                thinking_msg = await ws.receive_json()
                self.assertEqual(thinking_msg["event"], "thinking")

                # 6. Check response_text event
                resp_text_msg = await ws.receive_json()
                self.assertEqual(resp_text_msg["event"], "response_text")
                self.assertEqual(resp_text_msg["text"], "Echo: Hello world")

                # 7. Check audio_start event
                audio_start_msg = await ws.receive_json()
                self.assertEqual(audio_start_msg["event"], "audio_start")
                self.assertEqual(audio_start_msg["content_type"], "audio/wav")
                audio_bytes_len = audio_start_msg["bytes"]

                # 8. Check binary audio chunks
                received_audio = b""
                while True:
                    msg = await ws.receive()
                    if msg.type == web.WSMsgType.BINARY:
                        received_audio += msg.data
                    elif msg.type == web.WSMsgType.TEXT:
                        data = json.loads(msg.data)
                        if data["event"] == "audio_end":
                            break
                    else:
                        self.fail(f"Unexpected msg type: {msg.type}")

                self.assertEqual(len(received_audio), audio_bytes_len)

                # 9. Check done event
                done_msg = await ws.receive_json()
                self.assertEqual(done_msg["event"], "done")

    async def test_disconnect_cleanup(self):
        """Verify that raw files are cleaned up if client disconnects abruptly."""
        headers = {
            "Authorization": "Bearer valid_secret_key",
            "X-Device-ID": "disconnect-device"
        }
        
        # We need to find the file created during this session.
        # Let's verify files in ~/.hermes/cache/audio/esp32 or /tmp/esp32
        audio_dir = Path.home() / ".hermes" / "cache" / "audio" / "esp32"
        existing_files_before = set(audio_dir.glob("ws_disconnect_device_*.s16le")) if audio_dir.exists() else set()

        async with ClientSession() as session:
            async with session.ws_connect(self.ws_url, headers=headers) as ws:
                await ws.send_str(json.dumps({
                    "event": "start",
                    "device_id": "disconnect-device"
                }))
                
                await ws.receive_json() # ready
                await ws.receive_json() # listening
                
                # Send some bytes
                await ws.send_bytes(b"\x00" * 1000)
                
                # Abrupt close/exit without 'stop'
                await ws.close()

        # Allow loop to process close
        await asyncio.sleep(0.5)

        existing_files_after = set(audio_dir.glob("ws_disconnect_device_*.s16le")) if audio_dir.exists() else set()
        new_files = existing_files_after - existing_files_before
        
        # Verify no orphaned raw files remaining
        self.assertEqual(len(new_files), 0, f"Found orphaned temp files: {new_files}")

    async def test_too_short_audio(self):
        """Verify session is rejected with audio_too_short error if total PCM is below 3200 bytes."""
        headers = {
            "Authorization": "Bearer valid_secret_key",
            "X-Device-ID": "short-device"
        }
        async with ClientSession() as session:
            async with session.ws_connect(self.ws_url, headers=headers) as ws:
                await ws.send_str(json.dumps({
                    "event": "start",
                    "device_id": "short-device"
                }))
                await ws.receive_json() # ready
                await ws.receive_json() # listening
                
                # Send only 1000 bytes (less than 3200 bytes limit)
                await ws.send_bytes(b"\x00" * 1000)
                await ws.send_str(json.dumps({"event": "stop"}))
                
                error_msg = await ws.receive_json()
                self.assertEqual(error_msg["event"], "error")
                self.assertEqual(error_msg["code"], "audio_too_short")


if __name__ == "__main__":
    unittest.main()