import json import os from pathlib import Path import socket import struct import tempfile import threading import time import unittest from unittest.mock import Mock, patch from pi_dictate_service.server import ( MAX_CLIENT_CONNECTIONS, DictateServiceCore, inherited_listener, main, serve, ) class LoadedModel: loaded = True def load(self) -> None: return None def transcribe(self, samples: list[float]) -> str: return "" class UnixLoopTests(unittest.TestCase): def setUp(self) -> None: self.temporary = tempfile.TemporaryDirectory() self.socket_path = Path(self.temporary.name) / "dictate.sock" self.listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) self.listener.bind(str(self.socket_path)) self.listener.listen() self.core = DictateServiceCore(model=LoadedModel(), idle_seconds=0.4) self.failure: BaseException | None = None def run() -> None: try: serve(self.listener, self.core) except BaseException as error: self.failure = error self.thread = threading.Thread(target=run) self.thread.start() def tearDown(self) -> None: self.thread.join(timeout=2) if self.thread.is_alive(): self.fail("server loop did not exit after its idle timeout") if self.failure: raise self.failure self.temporary.cleanup() def connect(self) -> socket.socket: client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) client.settimeout(1) client.connect(str(self.socket_path)) return client @staticmethod def receive_frame(client: socket.socket) -> dict[str, object]: data = bytearray() while not data.endswith(b"\n"): chunk = client.recv(4096) if not chunk: break data.extend(chunk) return json.loads(data) def test_accepts_fragmented_requests_and_writes_persistent_responses(self) -> None: with self.connect() as client: client.sendall(b'{"version":2,"type":"hello",') client.sendall(b'"requestId":"hello"}\n') self.assertEqual(self.receive_frame(client)["type"], "pong") client.sendall(b'{"version":2,"type":"ping","requestId":"ping"}\n') pong = self.receive_frame(client) self.assertEqual(pong["requestId"], "ping") self.assertEqual(pong["service"], "pi-dictate") def test_valid_state_error_keeps_open_but_malformed_frame_errors_once_then_closes(self) -> None: with self.connect() as client: client.sendall(b'{"version":2,"type":"hello","requestId":"hello"}\n') self.receive_frame(client) client.sendall( b'{"version":2,"type":"stop","requestId":"stop","generationId":"g"}\n' ) self.assertEqual(self.receive_frame(client)["code"], "invalid_state") client.sendall(b'{"version":2,"type":"ping","requestId":"ping"}\n') self.assertEqual(self.receive_frame(client)["type"], "pong") client.sendall(b'{broken}\n') self.assertEqual(self.receive_frame(client)["code"], "malformed_json") self.assertEqual(client.recv(1), b"") self.assertTrue(self.socket_path.exists(), "the service must not unlink a listener-owned path") def test_malicious_json_values_do_not_crash_the_service(self) -> None: with self.connect() as client: client.sendall(b'{"version":2,"type":[],"requestId":"bad-type"}\n') self.assertEqual(self.receive_frame(client)["code"], "invalid_request") client.sendall(b'{"version":2,"type":"ping","requestId":"\\ud800"}\n') self.assertEqual(self.receive_frame(client)["code"], "invalid_request") with self.connect() as client: huge_integer = b"9" * 5000 client.sendall( b'{"version":' + huge_integer + b',"type":"hello","requestId":"huge"}\n' ) self.assertEqual(self.receive_frame(client)["code"], "malformed_json") self.assertEqual(client.recv(1), b"") with self.connect() as client: client.sendall(b'{"version":2,"type":"hello","requestId":"healthy"}\n') self.assertEqual(self.receive_frame(client)["type"], "pong") def test_closes_non_reading_client_before_output_queue_can_grow_unbounded(self) -> None: with self.connect() as client: client.sendall(b'{"version":2,"type":"hello","requestId":"hello"}\n') self.receive_frame(client) ping = b'{"version":2,"type":"ping","requestId":"p"}\n' try: client.sendall(ping * 1200) except BrokenPipeError: pass received = bytearray() while True: try: chunk = client.recv(64 * 1024) except ConnectionResetError: break if not chunk: break received.extend(chunk) self.assertLessEqual(len(received), 64 * 1024) with self.connect() as client: client.sendall(b'{"version":2,"type":"hello","requestId":"healthy"}\n') self.assertEqual(self.receive_frame(client)["type"], "pong") def test_rejects_connections_above_the_fixed_limit(self) -> None: clients = [self.connect() for _ in range(MAX_CLIENT_CONNECTIONS)] try: time.sleep(0.1) overflow = self.connect() try: overflow.sendall(b'{"version":2,"type":"hello","requestId":"overflow"}\n') with self.assertRaises((ConnectionResetError, BrokenPipeError)): if overflow.recv(1) == b"": raise ConnectionResetError finally: overflow.close() finally: for client in clients: client.close() def test_client_reset_does_not_crash_the_service(self) -> None: reset_client = self.connect() reset_client.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0)) reset_client.sendall(b'{"version":2,"type":"hello","requestId":"abandoned"}\n') reset_client.close() time.sleep(0.15) with self.connect() as client: client.sendall(b'{"version":2,"type":"hello","requestId":"hello"}\n') self.assertEqual(self.receive_frame(client)["type"], "pong") class ActivationTests(unittest.TestCase): def test_inherited_listener_requires_one_fd_for_this_process_and_wraps_fd_three(self) -> None: listener = Mock() with patch("pi_dictate_service.server.os.getpid", return_value=42), patch( "pi_dictate_service.server.socket.socket", return_value=listener ) as socket_factory: self.assertIs( inherited_listener({"LISTEN_PID": "42", "LISTEN_FDS": "1"}), listener, ) socket_factory.assert_called_once_with(fileno=3) listener.setblocking.assert_called_once_with(False) with patch("pi_dictate_service.server.os.getpid", return_value=42): with self.assertRaisesRegex(RuntimeError, "exactly one"): inherited_listener({"LISTEN_PID": "41", "LISTEN_FDS": "1"}) def test_main_requires_local_model_path_and_reads_idle_override(self) -> None: with patch.dict(os.environ, {}, clear=True): with self.assertRaisesRegex(SystemExit, "PI_DICTATE_MODEL_PATH"): main() model = object() listener = object() with patch.dict( os.environ, { "PI_DICTATE_MODEL_PATH": "/models/parakeet.nemo", "PI_DICTATE_TERMS_PATH": "/config/coding-terms.txt", "PI_DICTATE_BOOST_ALPHA": "0.5", "PI_DICTATE_IDLE_SECONDS": "12.5", }, clear=True, ), patch("pi_dictate_service.server.NemoParakeetModel", return_value=model) as adapter, patch( "pi_dictate_service.server.DictateServiceCore" ) as core_factory, patch( "pi_dictate_service.server.inherited_listener", return_value=listener ), patch("pi_dictate_service.server.serve") as serve_loop: core = core_factory.return_value main() adapter.assert_called_once_with( "/models/parakeet.nemo", terms_path="/config/coding-terms.txt", boost_alpha=0.5, ) core_factory.assert_called_once_with(model=model, idle_seconds=12.5) serve_loop.assert_called_once_with(listener, core) if __name__ == "__main__": unittest.main()