Implement WebSocket real-time audio streaming for Hermes gateway

hermes_voice_gateway/test_websocket.py

@@ -0,0 +1,230 @@
+# 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()