feat(dictate): use local socket-activated Parakeet
This commit is contained in:
@@ -0,0 +1,107 @@
|
||||
from pathlib import Path
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
|
||||
SCRIPT = Path(__file__).resolve().parents[1] / "install.sh"
|
||||
|
||||
|
||||
class InstallerTests(unittest.TestCase):
|
||||
def run_installer(
|
||||
self,
|
||||
home: Path,
|
||||
model: Path,
|
||||
*arguments: str,
|
||||
install_root_override: str | None = None,
|
||||
) -> subprocess.CompletedProcess[str]:
|
||||
environment = {
|
||||
**os.environ,
|
||||
"HOME": str(home),
|
||||
"XDG_CONFIG_HOME": str(home / ".config"),
|
||||
"PI_DICTATE_INSTALL_ROOT": install_root_override or str(
|
||||
home / ".local/share/pi-dictate"
|
||||
),
|
||||
}
|
||||
return subprocess.run(
|
||||
[str(SCRIPT), *arguments, "--dry-run", "--model-source", str(model)],
|
||||
env=environment,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
|
||||
def test_dry_run_seeds_dictionary_and_installs_owned_unit_copies(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
home = Path(directory)
|
||||
model = home / "model.nemo"
|
||||
model.write_bytes(b"model")
|
||||
|
||||
result = self.run_installer(home, model)
|
||||
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
self.assertIn("coding-terms.txt", result.stdout)
|
||||
self.assertIn("pi-dictate.socket", result.stdout)
|
||||
self.assertIn("install -m 0644", result.stdout)
|
||||
self.assertNotIn("ln -s", result.stdout)
|
||||
self.assertIn(".nemo.install.", result.stdout)
|
||||
self.assertIn("mv -f", result.stdout)
|
||||
self.assertIn("uv sync --frozen --no-dev --no-editable", result.stdout)
|
||||
self.assertIn("Would install pi-dictate", result.stdout)
|
||||
self.assertFalse((home / ".local/share/pi-dictate").exists())
|
||||
|
||||
def test_ignores_dangerous_install_root_environment_override(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
home = Path(directory)
|
||||
model = home / "model.nemo"
|
||||
model.write_bytes(b"model")
|
||||
|
||||
result = self.run_installer(
|
||||
home,
|
||||
model,
|
||||
"uninstall",
|
||||
install_root_override="/",
|
||||
)
|
||||
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
self.assertIn(str(home / ".local/share/pi-dictate"), result.stdout)
|
||||
self.assertNotIn("rm -rf / ", result.stdout)
|
||||
|
||||
def test_uninstall_refuses_to_disable_an_unmanaged_same_named_unit(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
home = Path(directory)
|
||||
model = home / "model.nemo"
|
||||
model.write_bytes(b"model")
|
||||
unit = home / ".config/systemd/user/pi-dictate.socket"
|
||||
unit.parent.mkdir(parents=True)
|
||||
unit.write_text("[Socket]\nListenStream=/tmp/unrelated.sock\n", encoding="utf-8")
|
||||
|
||||
result = self.run_installer(home, model, "uninstall")
|
||||
|
||||
self.assertEqual(result.returncode, 1)
|
||||
self.assertIn("unmanaged systemd unit", result.stderr)
|
||||
self.assertNotIn("systemctl --user disable", result.stdout)
|
||||
self.assertEqual(
|
||||
unit.read_text(encoding="utf-8"),
|
||||
"[Socket]\nListenStream=/tmp/unrelated.sock\n",
|
||||
)
|
||||
|
||||
def test_reinstall_preserves_an_existing_host_dictionary(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
home = Path(directory)
|
||||
model = home / "model.nemo"
|
||||
model.write_bytes(b"model")
|
||||
dictionary = home / ".local/share/pi-dictate/config/coding-terms.txt"
|
||||
dictionary.parent.mkdir(parents=True)
|
||||
dictionary.write_text("my private term\n", encoding="utf-8")
|
||||
|
||||
result = self.run_installer(home, model)
|
||||
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
self.assertIn("Keeping existing coding dictionary", result.stdout)
|
||||
self.assertEqual(dictionary.read_text(encoding="utf-8"), "my private term\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,143 @@
|
||||
import hashlib
|
||||
import os
|
||||
from contextlib import nullcontext
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
from types import ModuleType, SimpleNamespace
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from pi_dictate_service.model import NemoParakeetModel, load_coding_terms
|
||||
|
||||
|
||||
class FakeNemoModel:
|
||||
def __init__(self) -> None:
|
||||
self.eval_called = False
|
||||
self.to_device: str | None = None
|
||||
self.audio = None
|
||||
self.transcribe_options: dict[str, object] = {}
|
||||
self.cfg = SimpleNamespace(decoding={"strategy": "greedy_batch", "greedy": {"max_symbols": 10}})
|
||||
self.decoding_config: dict[str, object] | None = None
|
||||
|
||||
def eval(self) -> None:
|
||||
self.eval_called = True
|
||||
|
||||
def to(self, device: str) -> None:
|
||||
self.to_device = device
|
||||
|
||||
def change_decoding_strategy(self, config: dict[str, object], verbose: bool) -> None:
|
||||
self.decoding_config = config
|
||||
|
||||
def transcribe(self, audio: list[object], **options: object) -> list[object]:
|
||||
self.audio = audio
|
||||
self.transcribe_options = options
|
||||
return [type("Hypothesis", (), {"text": "hello locally"})()]
|
||||
|
||||
|
||||
class FakeASRModel:
|
||||
restored_path: str | None = None
|
||||
restored: FakeNemoModel | None = None
|
||||
|
||||
@classmethod
|
||||
def restore_from(cls, restore_path: str, map_location: str) -> FakeNemoModel:
|
||||
cls.restored_path = restore_path
|
||||
self = FakeNemoModel()
|
||||
cls.restored = self
|
||||
return self
|
||||
|
||||
|
||||
class FakeNumpy(ModuleType):
|
||||
float32 = "float32"
|
||||
|
||||
def asarray(self, samples: list[float], dtype: object) -> tuple[list[float], object]:
|
||||
return samples, dtype
|
||||
|
||||
|
||||
class NemoParakeetModelTests(unittest.TestCase):
|
||||
def test_load_is_lazy_offline_applies_coding_terms_and_transcribes_in_memory_audio(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
model_path = Path(directory) / "parakeet.nemo"
|
||||
model_path.write_bytes(b"model")
|
||||
terms_path = Path(directory) / "coding-terms.txt"
|
||||
terms_path.write_text("JSON\nnpm test\n", encoding="utf-8")
|
||||
modules = self.fake_modules()
|
||||
with patch.dict(sys.modules, modules, clear=False), patch.dict(os.environ, {}, clear=True):
|
||||
adapter = NemoParakeetModel(
|
||||
model_path,
|
||||
terms_path=terms_path,
|
||||
boost_alpha=0.5,
|
||||
expected_sha256=hashlib.sha256(b"model").hexdigest(),
|
||||
)
|
||||
self.assertIsNone(FakeASRModel.restored_path)
|
||||
|
||||
adapter.load()
|
||||
transcript = adapter.transcribe([0.0, 0.5])
|
||||
|
||||
self.assertEqual(os.environ["HF_HUB_OFFLINE"], "1")
|
||||
self.assertEqual(os.environ["TRANSFORMERS_OFFLINE"], "1")
|
||||
self.assertEqual(FakeASRModel.restored_path, str(model_path))
|
||||
self.assertTrue(FakeASRModel.restored.eval_called)
|
||||
self.assertEqual(FakeASRModel.restored.to_device, "cuda")
|
||||
self.assertEqual(transcript, "hello locally")
|
||||
self.assertEqual(FakeASRModel.restored.audio, [([0.0, 0.5], "float32")])
|
||||
self.assertEqual(
|
||||
FakeASRModel.restored.decoding_config["greedy"]["boosting_tree"]["key_phrases_list"],
|
||||
["JSON", "npm test"],
|
||||
)
|
||||
self.assertEqual(
|
||||
FakeASRModel.restored.decoding_config["greedy"]["boosting_tree_alpha"],
|
||||
0.5,
|
||||
)
|
||||
self.assertEqual(
|
||||
FakeASRModel.restored.transcribe_options,
|
||||
{"batch_size": 1, "num_workers": 0, "verbose": False},
|
||||
)
|
||||
|
||||
def test_loads_a_bounded_deduplicated_host_dictionary(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
terms_path = Path(directory) / "coding-terms.txt"
|
||||
terms_path.write_text("# local terms\n JSON \n\nnpm\nJSON\n", encoding="utf-8")
|
||||
self.assertEqual(load_coding_terms(terms_path), ["JSON", "npm"])
|
||||
|
||||
terms_path.write_text("x" * 129, encoding="utf-8")
|
||||
with self.assertRaisesRegex(ValueError, "128 UTF-8 bytes"):
|
||||
load_coding_terms(terms_path)
|
||||
|
||||
def test_refuses_a_model_that_no_longer_matches_the_pinned_hash(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
model_path = Path(directory) / "parakeet.nemo"
|
||||
model_path.write_bytes(b"replaced")
|
||||
adapter = NemoParakeetModel(model_path)
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "model checksum mismatch"):
|
||||
adapter.load()
|
||||
|
||||
def test_requires_an_existing_local_nemo_file(self) -> None:
|
||||
with self.assertRaisesRegex(ValueError, "local .nemo"):
|
||||
NemoParakeetModel("/does/not/exist")
|
||||
|
||||
@staticmethod
|
||||
def fake_modules() -> dict[str, ModuleType]:
|
||||
nemo = ModuleType("nemo")
|
||||
collections = ModuleType("nemo.collections")
|
||||
asr = ModuleType("nemo.collections.asr")
|
||||
models = ModuleType("nemo.collections.asr.models")
|
||||
models.ASRModel = FakeASRModel
|
||||
asr.models = models
|
||||
collections.asr = asr
|
||||
nemo.collections = collections
|
||||
omegaconf = ModuleType("omegaconf")
|
||||
omegaconf.open_dict = lambda _config: nullcontext()
|
||||
return {
|
||||
"nemo": nemo,
|
||||
"nemo.collections": collections,
|
||||
"nemo.collections.asr": asr,
|
||||
"nemo.collections.asr.models": models,
|
||||
"numpy": FakeNumpy("numpy"),
|
||||
"omegaconf": omegaconf,
|
||||
}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,162 @@
|
||||
import json
|
||||
import unittest
|
||||
|
||||
from pi_dictate_service.audio import AudioError, pcm_s16le_to_floats
|
||||
from pi_dictate_service.protocol import (
|
||||
MAX_ERROR_DETAIL_BYTES,
|
||||
MAX_FRAME_BYTES,
|
||||
MAX_ID_BYTES,
|
||||
MAX_TRANSCRIPT_BYTES,
|
||||
FrameDecoder,
|
||||
ProtocolError,
|
||||
encode_frame,
|
||||
parse_request,
|
||||
)
|
||||
|
||||
|
||||
class FrameDecoderTests(unittest.TestCase):
|
||||
def test_decodes_fragmented_and_multiple_frames(self) -> None:
|
||||
decoder = FrameDecoder()
|
||||
|
||||
self.assertEqual(decoder.feed(b'{"version":1'), [])
|
||||
frames = decoder.feed(
|
||||
b',"type":"ping","requestId":"one"}\n'
|
||||
b'{"version":1,"type":"ping","requestId":"two"}\n'
|
||||
)
|
||||
|
||||
self.assertEqual([frame["requestId"] for frame in frames], ["one", "two"])
|
||||
|
||||
def test_rejects_malformed_json_and_oversized_frames_as_fatal(self) -> None:
|
||||
with self.assertRaisesRegex(ProtocolError, "malformed JSON") as malformed:
|
||||
FrameDecoder().feed(b'{nope}\n')
|
||||
self.assertTrue(malformed.exception.close_connection)
|
||||
self.assertEqual(malformed.exception.code, "malformed_json")
|
||||
|
||||
with self.assertRaises(ProtocolError) as oversized:
|
||||
FrameDecoder().feed(b"x" * (MAX_FRAME_BYTES + 1))
|
||||
self.assertTrue(oversized.exception.close_connection)
|
||||
self.assertEqual(oversized.exception.code, "frame_too_large")
|
||||
|
||||
|
||||
class RequestValidationTests(unittest.TestCase):
|
||||
def test_requires_exact_fields_bounded_ids_and_supported_version(self) -> None:
|
||||
request = parse_request(
|
||||
{"version": 1, "type": "start", "requestId": "r1", "generationId": "g1"}
|
||||
)
|
||||
self.assertEqual(request.generation_id, "g1")
|
||||
|
||||
with self.assertRaises(ProtocolError) as unknown:
|
||||
parse_request(
|
||||
{
|
||||
"version": 1,
|
||||
"type": "start",
|
||||
"requestId": "r1",
|
||||
"generationId": "g1",
|
||||
"path": "/tmp/audio",
|
||||
}
|
||||
)
|
||||
self.assertEqual(unknown.exception.code, "invalid_request")
|
||||
self.assertFalse(unknown.exception.close_connection)
|
||||
|
||||
with self.assertRaises(ProtocolError) as long_id:
|
||||
parse_request(
|
||||
{
|
||||
"version": 1,
|
||||
"type": "ping",
|
||||
"requestId": "é" * ((MAX_ID_BYTES // 2) + 1),
|
||||
}
|
||||
)
|
||||
self.assertEqual(long_id.exception.code, "invalid_request")
|
||||
|
||||
with self.assertRaises(ProtocolError) as unsupported:
|
||||
parse_request({"version": 2, "type": "hello", "requestId": "r2"})
|
||||
self.assertEqual(unsupported.exception.code, "unsupported_version")
|
||||
self.assertTrue(unsupported.exception.close_connection)
|
||||
|
||||
def test_rejects_unhashable_types_invalid_unicode_and_huge_integers(self) -> None:
|
||||
for request_type in ([], {}):
|
||||
with self.subTest(request_type=request_type):
|
||||
with self.assertRaises(ProtocolError) as failed:
|
||||
parse_request(
|
||||
{"version": 1, "type": request_type, "requestId": "request"}
|
||||
)
|
||||
self.assertEqual(failed.exception.code, "invalid_request")
|
||||
|
||||
with self.assertRaises(ProtocolError) as surrogate:
|
||||
parse_request({"version": 1, "type": "ping", "requestId": "\ud800"})
|
||||
self.assertEqual(surrogate.exception.request_id, "")
|
||||
|
||||
huge_integer = b'{"version":' + (b"9" * 5000) + b',"type":"hello","requestId":"r"}\n'
|
||||
with self.assertRaises(ProtocolError) as huge:
|
||||
FrameDecoder().feed(huge_integer)
|
||||
self.assertEqual(huge.exception.code, "malformed_json")
|
||||
self.assertTrue(huge.exception.close_connection)
|
||||
|
||||
def test_error_echoes_only_valid_bounded_ids(self) -> None:
|
||||
oversized = "x" * (MAX_ID_BYTES + 1)
|
||||
with self.assertRaises(ProtocolError) as unsupported:
|
||||
parse_request(
|
||||
{"version": 2, "type": "hello", "requestId": oversized}
|
||||
)
|
||||
self.assertEqual(unsupported.exception.request_id, "")
|
||||
|
||||
encoded = encode_frame(
|
||||
{
|
||||
"version": 1,
|
||||
"type": "error",
|
||||
"requestId": oversized,
|
||||
"generationId": "\ud800",
|
||||
"code": "invalid_request",
|
||||
"detail": "bad request",
|
||||
}
|
||||
)
|
||||
payload = json.loads(encoded)
|
||||
self.assertEqual(payload["requestId"], "")
|
||||
self.assertNotIn("generationId", payload)
|
||||
|
||||
def test_encoded_transcripts_and_error_details_are_bounded(self) -> None:
|
||||
with self.assertRaises(ProtocolError) as transcript:
|
||||
encode_frame(
|
||||
{
|
||||
"version": 1,
|
||||
"type": "result",
|
||||
"generationId": "g1",
|
||||
"transcript": "é" * ((MAX_TRANSCRIPT_BYTES // 2) + 1),
|
||||
}
|
||||
)
|
||||
self.assertEqual(transcript.exception.code, "transcript_too_large")
|
||||
|
||||
encoded = encode_frame(
|
||||
{
|
||||
"version": 1,
|
||||
"type": "error",
|
||||
"requestId": "r1",
|
||||
"code": "failed",
|
||||
"detail": "é" * MAX_ERROR_DETAIL_BYTES,
|
||||
}
|
||||
)
|
||||
payload = json.loads(encoded)
|
||||
self.assertLessEqual(len(payload["detail"].encode("utf-8")), MAX_ERROR_DETAIL_BYTES)
|
||||
self.assertLessEqual(len(encoded) - 1, MAX_FRAME_BYTES)
|
||||
|
||||
|
||||
class AudioConversionTests(unittest.TestCase):
|
||||
def test_converts_little_endian_signed_16_bit_pcm(self) -> None:
|
||||
pcm = b"\x00\x80\x00\xc0\x00\x00\x00\x40\xff\x7f"
|
||||
|
||||
samples = pcm_s16le_to_floats(pcm)
|
||||
self.assertEqual(samples.itemsize, 4)
|
||||
self.assertEqual(
|
||||
list(samples),
|
||||
[-1.0, -0.5, 0.0, 0.5, 32767 / 32768.0],
|
||||
)
|
||||
|
||||
def test_rejects_odd_or_overlong_pcm(self) -> None:
|
||||
with self.assertRaisesRegex(AudioError, "even"):
|
||||
pcm_s16le_to_floats(b"\x00")
|
||||
with self.assertRaisesRegex(AudioError, "duration"):
|
||||
pcm_s16le_to_floats(b"\x00\x00" * 5, max_bytes=8)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,246 @@
|
||||
import io
|
||||
import subprocess
|
||||
import threading
|
||||
import unittest
|
||||
|
||||
from pi_dictate_service.recorder import PW_RECORD_COMMAND, PwRecorder, RecorderError
|
||||
|
||||
|
||||
class TrackingStream(io.BytesIO):
|
||||
def __init__(self, value: bytes, events: list[str], name: str) -> None:
|
||||
super().__init__(value)
|
||||
self._events = events
|
||||
self._name = name
|
||||
self.eof = threading.Event()
|
||||
|
||||
def read(self, size: int = -1) -> bytes:
|
||||
value = super().read(size)
|
||||
if not value and not self.eof.is_set():
|
||||
self._events.append(f"{self._name}-eof")
|
||||
self.eof.set()
|
||||
return value
|
||||
|
||||
|
||||
class FakeProcess:
|
||||
def __init__(
|
||||
self,
|
||||
stdout: TrackingStream,
|
||||
stderr: TrackingStream,
|
||||
*,
|
||||
returncode: int | None = None,
|
||||
timeout: bool = False,
|
||||
terminate_returncode: int = -15,
|
||||
) -> None:
|
||||
self.pid = 1234
|
||||
self.stdout = stdout
|
||||
self.stderr = stderr
|
||||
self.returncode = returncode
|
||||
self._timeout = timeout
|
||||
self._terminate_returncode = terminate_returncode
|
||||
self.events = stdout._events
|
||||
|
||||
def poll(self) -> int | None:
|
||||
return self.returncode
|
||||
|
||||
def terminate(self) -> None:
|
||||
self.events.append("terminate")
|
||||
if not self._timeout:
|
||||
self.returncode = self._terminate_returncode
|
||||
|
||||
def kill(self) -> None:
|
||||
self.events.append("kill")
|
||||
self.returncode = -9
|
||||
self._timeout = False
|
||||
|
||||
def wait(self, timeout: float | None = None) -> int:
|
||||
self.events.append("wait")
|
||||
if self._timeout:
|
||||
raise subprocess.TimeoutExpired("pw-record", timeout)
|
||||
if self.returncode is None:
|
||||
self.returncode = 0
|
||||
return self.returncode
|
||||
|
||||
|
||||
class ProcessFactory:
|
||||
def __init__(self, process: FakeProcess) -> None:
|
||||
self.process = process
|
||||
self.calls: list[tuple[list[str], dict[str, object]]] = []
|
||||
|
||||
def __call__(self, command: list[str], **kwargs: object) -> FakeProcess:
|
||||
self.calls.append((command, kwargs))
|
||||
return self.process
|
||||
|
||||
|
||||
class Clock:
|
||||
def __init__(self) -> None:
|
||||
self.now = 0.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
|
||||
class PwRecorderTests(unittest.TestCase):
|
||||
def make_recorder(
|
||||
self,
|
||||
pcm: bytes = b"\x00\x00\x01\x00",
|
||||
stderr: bytes = b"",
|
||||
*,
|
||||
returncode: int | None = None,
|
||||
timeout: bool = False,
|
||||
max_bytes: int = 1024,
|
||||
stderr_limit: int = 64 * 1024,
|
||||
clock: Clock | None = None,
|
||||
terminate_returncode: int = -15,
|
||||
) -> tuple[PwRecorder, FakeProcess, ProcessFactory, list[str]]:
|
||||
events: list[str] = []
|
||||
process = FakeProcess(
|
||||
TrackingStream(pcm, events, "stdout"),
|
||||
TrackingStream(stderr, events, "stderr"),
|
||||
returncode=returncode,
|
||||
timeout=timeout,
|
||||
terminate_returncode=terminate_returncode,
|
||||
)
|
||||
factory = ProcessFactory(process)
|
||||
def signal_group(pid: int, signal_number: int) -> None:
|
||||
self.assertEqual(pid, process.pid)
|
||||
if signal_number == 15:
|
||||
process.terminate()
|
||||
elif signal_number == 9:
|
||||
process.kill()
|
||||
else:
|
||||
self.fail(f"unexpected signal: {signal_number}")
|
||||
|
||||
recorder = PwRecorder(
|
||||
process_factory=factory,
|
||||
max_bytes=max_bytes,
|
||||
stderr_limit=stderr_limit,
|
||||
stop_timeout=0.01,
|
||||
clock=clock or Clock(),
|
||||
process_group_signal=signal_group,
|
||||
)
|
||||
return recorder, process, factory, events
|
||||
|
||||
def test_spawns_exact_pw_record_command_and_drains_bounded_streams(self) -> None:
|
||||
recorder, process, factory, events = self.make_recorder(
|
||||
pcm=b"\x00\x00\x01\x00", stderr=b"0123456789", stderr_limit=8
|
||||
)
|
||||
|
||||
recorder.start()
|
||||
self.assertTrue(process.stdout.eof.wait(1))
|
||||
capture = recorder.stop()
|
||||
|
||||
self.assertEqual(factory.calls[0][0], PW_RECORD_COMMAND)
|
||||
self.assertIs(factory.calls[0][1]["stdout"], subprocess.PIPE)
|
||||
self.assertIs(factory.calls[0][1]["stderr"], subprocess.PIPE)
|
||||
self.assertIs(factory.calls[0][1]["start_new_session"], True)
|
||||
self.assertEqual(capture.pcm, b"\x00\x00\x01\x00")
|
||||
self.assertEqual(capture.stderr_tail, b"23456789")
|
||||
self.assertIn("stdout-eof", events)
|
||||
self.assertIn("wait", events)
|
||||
|
||||
def test_accepts_pw_record_code_one_after_intentional_sigterm(self) -> None:
|
||||
recorder, process, _factory, _events = self.make_recorder(terminate_returncode=1)
|
||||
recorder.start()
|
||||
|
||||
capture = recorder.stop()
|
||||
|
||||
self.assertEqual(capture.returncode, 1)
|
||||
self.assertIsNone(capture.signal)
|
||||
self.assertEqual(capture.pcm, b"\x00\x00\x01\x00")
|
||||
self.assertEqual(process.returncode, 1)
|
||||
|
||||
def test_rejects_code_one_if_process_already_exited_or_has_stderr(self) -> None:
|
||||
exited, exited_process, _factory, _events = self.make_recorder(returncode=1)
|
||||
exited.start()
|
||||
self.assertTrue(exited_process.stdout.eof.wait(1))
|
||||
with self.assertRaises(RecorderError):
|
||||
exited.stop()
|
||||
|
||||
noisy, _process, _factory, _events = self.make_recorder(
|
||||
stderr=b"pipewire failed",
|
||||
terminate_returncode=1,
|
||||
)
|
||||
noisy.start()
|
||||
with self.assertRaises(RecorderError):
|
||||
noisy.stop()
|
||||
|
||||
def test_shutdown_timeout_escalates_to_kill_and_fails(self) -> None:
|
||||
recorder, _process, _factory, events = self.make_recorder(timeout=True)
|
||||
recorder.start()
|
||||
|
||||
with self.assertRaises(RecorderError) as failed:
|
||||
recorder.stop()
|
||||
|
||||
self.assertEqual(failed.exception.code, "recorder_timeout")
|
||||
self.assertLess(events.index("terminate"), events.index("kill"))
|
||||
|
||||
def test_classifies_unexpected_exit_code_and_signal(self) -> None:
|
||||
recorder, process, _factory, _events = self.make_recorder(
|
||||
stderr=b"pipewire failed", returncode=7
|
||||
)
|
||||
recorder.start()
|
||||
self.assertTrue(process.stdout.eof.wait(1))
|
||||
|
||||
failure = recorder.poll_unexpected()
|
||||
|
||||
self.assertIsNotNone(failure)
|
||||
assert failure is not None
|
||||
self.assertEqual(failure.code, "recorder_exit")
|
||||
self.assertIn("code 7", failure.detail)
|
||||
self.assertIn("pipewire failed", failure.detail)
|
||||
|
||||
signalled, signal_process, _factory, _events = self.make_recorder(returncode=-11)
|
||||
signalled.start()
|
||||
self.assertTrue(signal_process.stdout.eof.wait(1))
|
||||
signal_failure = signalled.poll_unexpected()
|
||||
assert signal_failure is not None
|
||||
self.assertIn("signal 11", signal_failure.detail)
|
||||
|
||||
def test_duration_limit_discards_audio_and_reports_failure(self) -> None:
|
||||
recorder, process, _factory, _events = self.make_recorder(
|
||||
pcm=b"\x00\x00" * 10, max_bytes=8
|
||||
)
|
||||
recorder.start()
|
||||
self.assertTrue(process.stdout.eof.wait(1))
|
||||
|
||||
failure = recorder.poll_unexpected()
|
||||
|
||||
assert failure is not None
|
||||
self.assertEqual(failure.code, "duration_limit")
|
||||
self.assertEqual(recorder.buffered_bytes, 0)
|
||||
|
||||
def test_stop_reports_a_duration_limit_that_was_already_reached(self) -> None:
|
||||
recorder, process, _factory, _events = self.make_recorder(
|
||||
pcm=b"\x00\x00" * 10, max_bytes=8
|
||||
)
|
||||
recorder.start()
|
||||
self.assertTrue(process.stdout.eof.wait(1))
|
||||
|
||||
with self.assertRaises(RecorderError) as failed:
|
||||
recorder.stop()
|
||||
|
||||
self.assertEqual(failed.exception.code, "duration_limit")
|
||||
self.assertEqual(recorder.buffered_bytes, 0)
|
||||
|
||||
def test_duration_limit_kills_a_recorder_that_ignores_sigterm(self) -> None:
|
||||
clock = Clock()
|
||||
recorder, process, _factory, events = self.make_recorder(
|
||||
pcm=b"\x00\x00" * 10,
|
||||
max_bytes=8,
|
||||
timeout=True,
|
||||
clock=clock,
|
||||
)
|
||||
recorder.start()
|
||||
self.assertTrue(process.stdout.eof.wait(1))
|
||||
self.assertIsNone(recorder.poll_unexpected())
|
||||
|
||||
clock.now = 0.02
|
||||
failure = recorder.poll_unexpected()
|
||||
|
||||
assert failure is not None
|
||||
self.assertEqual(failure.code, "duration_limit")
|
||||
self.assertIn("kill", events)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,333 @@
|
||||
from concurrent.futures import Future
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
import unittest
|
||||
|
||||
from pi_dictate_service.protocol import MAX_TRANSCRIPT_BYTES, parse_request
|
||||
from pi_dictate_service.recorder import Capture, RecorderError
|
||||
from pi_dictate_service.server import DictateServiceCore, Outbound
|
||||
|
||||
|
||||
class Clock:
|
||||
def __init__(self) -> None:
|
||||
self.now = 0.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
def advance(self, seconds: float) -> None:
|
||||
self.now += seconds
|
||||
|
||||
|
||||
@dataclass
|
||||
class PendingWork:
|
||||
future: Future[object]
|
||||
function: object
|
||||
args: tuple[object, ...]
|
||||
|
||||
|
||||
class ControlledExecutor:
|
||||
def __init__(self) -> None:
|
||||
self.pending: list[PendingWork] = []
|
||||
|
||||
def submit(self, function: object, *args: object) -> Future[object]:
|
||||
future: Future[object] = Future()
|
||||
self.pending.append(PendingWork(future, function, args))
|
||||
return future
|
||||
|
||||
def run_next(self) -> None:
|
||||
work = self.pending.pop(0)
|
||||
try:
|
||||
result = work.function(*work.args) # type: ignore[operator]
|
||||
except BaseException as error:
|
||||
work.future.set_exception(error)
|
||||
else:
|
||||
work.future.set_result(result)
|
||||
|
||||
def shutdown(self, wait: bool = True, cancel_futures: bool = False) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class FakeModel:
|
||||
def __init__(self) -> None:
|
||||
self.loaded = False
|
||||
self.transcript = "dictated text"
|
||||
self.load_error: BaseException | None = None
|
||||
self.transcribe_error: BaseException | None = None
|
||||
self.samples: Sequence[float] | None = None
|
||||
|
||||
def load(self) -> None:
|
||||
if self.load_error:
|
||||
raise self.load_error
|
||||
self.loaded = True
|
||||
|
||||
def transcribe(self, samples: Sequence[float]) -> str:
|
||||
self.samples = samples
|
||||
if self.transcribe_error:
|
||||
raise self.transcribe_error
|
||||
return self.transcript
|
||||
|
||||
|
||||
class FakeRecorder:
|
||||
def __init__(self) -> None:
|
||||
self.started = False
|
||||
self.cancelled = False
|
||||
self.capture = Capture(b"\x00\x00\x00\x40", -15, 15, b"")
|
||||
self.stop_error: RecorderError | None = None
|
||||
self.unexpected: RecorderError | None = None
|
||||
|
||||
def start(self) -> None:
|
||||
self.started = True
|
||||
|
||||
def stop(self) -> Capture:
|
||||
if self.stop_error:
|
||||
raise self.stop_error
|
||||
return self.capture
|
||||
|
||||
def cancel(self) -> None:
|
||||
self.cancelled = True
|
||||
|
||||
def poll_unexpected(self) -> RecorderError | None:
|
||||
failure, self.unexpected = self.unexpected, None
|
||||
return failure
|
||||
|
||||
|
||||
class RecorderFactory:
|
||||
def __init__(self) -> None:
|
||||
self.instances: list[FakeRecorder] = []
|
||||
|
||||
def __call__(self) -> FakeRecorder:
|
||||
recorder = FakeRecorder()
|
||||
self.instances.append(recorder)
|
||||
return recorder
|
||||
|
||||
|
||||
class CoreTestCase(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.clock = Clock()
|
||||
self.executor = ControlledExecutor()
|
||||
self.model = FakeModel()
|
||||
self.recorders = RecorderFactory()
|
||||
self.core = DictateServiceCore(
|
||||
model=self.model,
|
||||
recorder_factory=self.recorders,
|
||||
executor=self.executor,
|
||||
clock=self.clock,
|
||||
idle_seconds=900,
|
||||
)
|
||||
|
||||
def request(
|
||||
self,
|
||||
client: str,
|
||||
request_type: str,
|
||||
request_id: str,
|
||||
generation_id: str | None = None,
|
||||
) -> list[dict[str, object]]:
|
||||
frame: dict[str, object] = {
|
||||
"version": 1,
|
||||
"type": request_type,
|
||||
"requestId": request_id,
|
||||
}
|
||||
if generation_id is not None:
|
||||
frame["generationId"] = generation_id
|
||||
self.core.handle_request(client, parse_request(frame))
|
||||
return self.frames(self.core.drain_events(), client)
|
||||
|
||||
@staticmethod
|
||||
def frames(events: list[Outbound], client: str | None = None) -> list[dict[str, object]]:
|
||||
return [event.frame for event in events if client is None or event.client_id == client]
|
||||
|
||||
def hello(self, client: str = "a") -> None:
|
||||
events = self.request(client, "hello", f"hello-{client}")
|
||||
self.assertEqual(events[0]["type"], "pong")
|
||||
|
||||
def load_to_listening(self, client: str = "a", generation: str = "g1") -> FakeRecorder:
|
||||
self.hello(client)
|
||||
self.request(client, "start", "start", generation)
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
events = self.frames(self.core.drain_events(), client)
|
||||
self.assertEqual(events[-1]["state"], "listening")
|
||||
self.assertEqual(events[-1]["requestId"], "start")
|
||||
return self.recorders.instances[-1]
|
||||
|
||||
|
||||
class HandshakeAndOwnershipTests(CoreTestCase):
|
||||
def test_hello_is_required_and_valid_errors_keep_session_open(self) -> None:
|
||||
events = self.request("a", "start", "s1", "g1")
|
||||
self.assertEqual(events[0]["code"], "hello_required")
|
||||
|
||||
self.hello("a")
|
||||
self.request("a", "start", "s2", "g1")
|
||||
self.assertEqual(len(self.executor.pending), 1)
|
||||
|
||||
def test_one_generation_is_globally_owned(self) -> None:
|
||||
self.load_to_listening()
|
||||
self.hello("b")
|
||||
|
||||
busy = self.request("b", "start", "s2", "g2")
|
||||
self.assertEqual(busy[0]["code"], "busy")
|
||||
not_owner = self.request("b", "stop", "x", "g1")
|
||||
self.assertEqual(not_owner[0]["code"], "not_owner")
|
||||
stale = self.request("a", "stop", "x2", "old")
|
||||
self.assertEqual(stale[0]["code"], "generation_mismatch")
|
||||
|
||||
stopping = self.request("a", "stop", "x3", "g1")
|
||||
self.assertEqual(stopping[0]["state"], "stopping")
|
||||
|
||||
|
||||
class CancellationTests(CoreTestCase):
|
||||
def test_cancel_during_load_is_immediate_and_blocks_until_worker_finishes(self) -> None:
|
||||
self.hello("a")
|
||||
self.request("a", "start", "start", "g1")
|
||||
|
||||
cancelled = self.request("a", "cancel", "cancel", "g1")
|
||||
self.assertEqual(cancelled[0]["type"], "cancelled")
|
||||
self.hello("b")
|
||||
self.assertEqual(self.request("b", "start", "start-b", "g2")[0]["code"], "busy")
|
||||
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
self.assertEqual(self.recorders.instances, [])
|
||||
self.assertNotIn("result", [event["type"] for event in self.frames(self.core.drain_events())])
|
||||
|
||||
def test_cancel_during_inference_suppresses_later_result(self) -> None:
|
||||
self.load_to_listening()
|
||||
self.request("a", "stop", "stop", "g1")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
transcribing = self.frames(self.core.drain_events(), "a")
|
||||
self.assertEqual(transcribing[0]["state"], "transcribing")
|
||||
|
||||
cancelled = self.request("a", "cancel", "cancel", "g1")
|
||||
self.assertEqual(cancelled[0]["type"], "cancelled")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
self.assertNotIn("result", [event["type"] for event in self.frames(self.core.drain_events())])
|
||||
|
||||
def test_disconnect_cancels_capture_and_tombstones_workers(self) -> None:
|
||||
recorder = self.load_to_listening()
|
||||
|
||||
self.core.disconnect("a")
|
||||
self.assertEqual(len(self.executor.pending), 1)
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
|
||||
self.assertTrue(recorder.cancelled)
|
||||
self.assertEqual(self.core.active_generation_id, None)
|
||||
|
||||
|
||||
class StopAndInferenceTests(CoreTestCase):
|
||||
def test_stop_drains_then_converts_before_inference_and_emits_result(self) -> None:
|
||||
self.load_to_listening()
|
||||
|
||||
stopping = self.request("a", "stop", "stop", "g1")
|
||||
self.assertEqual(stopping[0]["state"], "stopping")
|
||||
self.assertEqual(stopping[0]["requestId"], "stop")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
transcribing = self.frames(self.core.drain_events(), "a")
|
||||
self.assertEqual(transcribing[0]["state"], "transcribing")
|
||||
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
completed = self.frames(self.core.drain_events(), "a")
|
||||
|
||||
self.assertEqual(list(self.model.samples or []), [0.0, 0.5])
|
||||
self.assertEqual(completed[0]["type"], "result")
|
||||
self.assertEqual(completed[0]["transcript"], "dictated text")
|
||||
self.assertEqual(completed[1]["state"], "idle")
|
||||
|
||||
def test_empty_transcript_is_a_result(self) -> None:
|
||||
self.model.transcript = ""
|
||||
self.load_to_listening()
|
||||
self.request("a", "stop", "stop", "g1")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
self.core.drain_events()
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
|
||||
completed = self.frames(self.core.drain_events(), "a")
|
||||
self.assertEqual(completed[0]["type"], "result")
|
||||
self.assertEqual(completed[0]["transcript"], "")
|
||||
|
||||
def test_model_and_transcript_failures_are_bounded_terminal_errors(self) -> None:
|
||||
self.model.load_error = RuntimeError("CUDA unavailable")
|
||||
self.hello()
|
||||
self.request("a", "start", "start", "g1")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
failed = self.frames(self.core.drain_events(), "a")
|
||||
self.assertEqual(failed[0]["code"], "model_failure")
|
||||
self.assertEqual(failed[-1]["state"], "idle")
|
||||
|
||||
model = FakeModel()
|
||||
model.loaded = True
|
||||
model.transcript = "x" * (MAX_TRANSCRIPT_BYTES + 1)
|
||||
core = DictateServiceCore(
|
||||
model=model,
|
||||
recorder_factory=self.recorders,
|
||||
executor=self.executor,
|
||||
clock=self.clock,
|
||||
)
|
||||
core.handle_request("a", parse_request({"version": 1, "type": "hello", "requestId": "h"}))
|
||||
core.drain_events()
|
||||
core.handle_request("a", parse_request({"version": 1, "type": "start", "requestId": "s", "generationId": "g"}))
|
||||
core.drain_events()
|
||||
core.handle_request("a", parse_request({"version": 1, "type": "stop", "requestId": "t", "generationId": "g"}))
|
||||
self.executor.run_next()
|
||||
core.poll()
|
||||
core.drain_events()
|
||||
self.executor.run_next()
|
||||
core.poll()
|
||||
failed = self.frames(core.drain_events(), "a")
|
||||
self.assertEqual(failed[0]["code"], "transcript_too_large")
|
||||
|
||||
def test_unexpected_recorder_exit_is_terminal(self) -> None:
|
||||
recorder = self.load_to_listening()
|
||||
recorder.unexpected = RecorderError("recorder_exit", "pw-record exited with code 3")
|
||||
|
||||
self.core.poll()
|
||||
failed = self.frames(self.core.drain_events(), "a")
|
||||
|
||||
self.assertEqual(failed[0]["code"], "recorder_exit")
|
||||
self.assertEqual(failed[-1]["state"], "idle")
|
||||
|
||||
def test_stop_does_not_reclassify_an_already_exited_recorder_as_intentional(self) -> None:
|
||||
recorder = self.load_to_listening()
|
||||
recorder.unexpected = RecorderError("recorder_exit", "pw-record exited with code 0")
|
||||
|
||||
failed = self.request("a", "stop", "stop", "g1")
|
||||
|
||||
self.assertEqual(failed[0]["code"], "recorder_exit")
|
||||
self.assertEqual(failed[-1]["state"], "idle")
|
||||
self.assertEqual(self.executor.pending, [])
|
||||
|
||||
|
||||
class IdleTests(CoreTestCase):
|
||||
def test_only_valid_requests_reset_idle_and_workers_suspend_then_restart_timer(self) -> None:
|
||||
self.core.accepted("a")
|
||||
self.clock.advance(899)
|
||||
self.assertFalse(self.core.should_exit())
|
||||
self.request("a", "hello", "h")
|
||||
self.clock.advance(899)
|
||||
self.assertFalse(self.core.should_exit())
|
||||
|
||||
self.request("a", "start", "s", "g")
|
||||
self.clock.advance(901)
|
||||
self.assertFalse(self.core.should_exit())
|
||||
self.request("a", "cancel", "c", "g")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
self.core.drain_events()
|
||||
self.assertFalse(self.core.should_exit())
|
||||
|
||||
self.clock.advance(899)
|
||||
self.assertFalse(self.core.should_exit())
|
||||
self.clock.advance(1)
|
||||
self.assertTrue(self.core.should_exit())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,38 @@
|
||||
from pathlib import Path
|
||||
import unittest
|
||||
|
||||
|
||||
REPOSITORY_ROOT = Path(__file__).resolve().parents[3]
|
||||
UNIT_DIR = REPOSITORY_ROOT / "systemd/user"
|
||||
|
||||
|
||||
class SystemdUnitTests(unittest.TestCase):
|
||||
def test_socket_unit_owns_the_exact_private_runtime_path(self) -> None:
|
||||
socket_unit = (UNIT_DIR / "pi-dictate.socket").read_text(encoding="utf-8")
|
||||
self.assertIn("ListenStream=%t/pi-dictate.sock", socket_unit)
|
||||
self.assertIn("SocketMode=0600", socket_unit)
|
||||
self.assertIn("Accept=no", socket_unit)
|
||||
self.assertIn("RemoveOnStop=yes", socket_unit)
|
||||
|
||||
def test_service_is_offline_hardened_and_uses_private_native_caches(self) -> None:
|
||||
service_unit = (UNIT_DIR / "pi-dictate.service").read_text(encoding="utf-8")
|
||||
for expected in (
|
||||
"PI_DICTATE_MODEL_PATH=%h/.local/share/pi-dictate/models/",
|
||||
"PI_DICTATE_TERMS_PATH=%h/.local/share/pi-dictate/config/coding-terms.txt",
|
||||
"PI_DICTATE_BOOST_ALPHA=1.0",
|
||||
"HF_HUB_OFFLINE=1",
|
||||
"TRANSFORMERS_OFFLINE=1",
|
||||
"TRITON_CACHE_DIR=/tmp/pi-dictate-triton",
|
||||
"NUMBA_CACHE_DIR=/tmp/pi-dictate-numba",
|
||||
"LimitCORE=0",
|
||||
"MemorySwapMax=0",
|
||||
"ProtectSystem=strict",
|
||||
"ProtectHome=read-only",
|
||||
"RestrictAddressFamilies=AF_UNIX",
|
||||
"IPAddressDeny=any",
|
||||
):
|
||||
self.assertIn(expected, service_unit)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,225 @@
|
||||
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":1,"type":"hello",')
|
||||
client.sendall(b'"requestId":"hello"}\n')
|
||||
self.assertEqual(self.receive_frame(client)["type"], "pong")
|
||||
|
||||
client.sendall(b'{"version":1,"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":1,"type":"hello","requestId":"hello"}\n')
|
||||
self.receive_frame(client)
|
||||
client.sendall(
|
||||
b'{"version":1,"type":"stop","requestId":"stop","generationId":"g"}\n'
|
||||
)
|
||||
self.assertEqual(self.receive_frame(client)["code"], "invalid_state")
|
||||
client.sendall(b'{"version":1,"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":1,"type":[],"requestId":"bad-type"}\n')
|
||||
self.assertEqual(self.receive_frame(client)["code"], "invalid_request")
|
||||
client.sendall(b'{"version":1,"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":1,"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":1,"type":"hello","requestId":"hello"}\n')
|
||||
self.receive_frame(client)
|
||||
ping = b'{"version":1,"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":1,"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":1,"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":1,"type":"hello","requestId":"abandoned"}\n')
|
||||
reset_client.close()
|
||||
time.sleep(0.15)
|
||||
|
||||
with self.connect() as client:
|
||||
client.sendall(b'{"version":1,"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()
|
||||
Reference in New Issue
Block a user