226 lines
8.6 KiB
Python
226 lines
8.6 KiB
Python
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()
|