917 lines
29 KiB
Python
917 lines
29 KiB
Python
from concurrent.futures import Future
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass
|
|
import errno
|
|
import selectors
|
|
import subprocess
|
|
import sys
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from pi_dictate_service.protocol import (
|
|
MAX_FRAME_BYTES,
|
|
MAX_TRANSCRIPT_BYTES,
|
|
PROTOCOL_VERSION,
|
|
encode_frame,
|
|
parse_request,
|
|
)
|
|
from pi_dictate_service.recorder import Capture, RecorderError
|
|
from pi_dictate_service.server import DictateServiceCore, Outbound, serve
|
|
|
|
|
|
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.audio_level = 0.0
|
|
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 SpawnFailRecorder(FakeRecorder):
|
|
cleanup_complete = True
|
|
|
|
def start(self) -> None:
|
|
raise RecorderError("recorder_spawn", "could not start pw-record")
|
|
|
|
|
|
class SpawnFailRecorderFactory:
|
|
def __call__(self) -> SpawnFailRecorder:
|
|
return SpawnFailRecorder()
|
|
|
|
|
|
class DelayedCleanupRecorder(FakeRecorder):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.cleanup_complete = False
|
|
self.cleanup_attempts = 0
|
|
self.cleanup_started = False
|
|
|
|
def stop(self) -> Capture:
|
|
self.cleanup_started = True
|
|
raise RecorderError("recorder_timeout", "pw-record cleanup did not complete")
|
|
|
|
def poll_unexpected(self) -> RecorderError | None:
|
|
if not self.cleanup_started:
|
|
return None
|
|
self.cleanup_attempts += 1
|
|
if self.cleanup_attempts >= 3:
|
|
self.cleanup_complete = True
|
|
return None
|
|
|
|
|
|
class DelayedCleanupRecorderFactory:
|
|
def __init__(self) -> None:
|
|
self.instance = DelayedCleanupRecorder()
|
|
|
|
def __call__(self) -> DelayedCleanupRecorder:
|
|
return self.instance
|
|
|
|
|
|
class DelayedCloseTurn:
|
|
sample_rate_hz = 16_000
|
|
retain_pcm = True
|
|
|
|
def __init__(self) -> None:
|
|
self.close_calls = 0
|
|
|
|
@property
|
|
def cleanup_complete(self) -> bool:
|
|
return self.close_calls >= 3
|
|
|
|
def enqueue_pcm(self, chunk: bytes) -> None:
|
|
return None
|
|
|
|
def finish(self, capture: Capture | None = None) -> str:
|
|
return "cloud result"
|
|
|
|
def poll_error(self) -> None:
|
|
return None
|
|
|
|
def cancel(self) -> None:
|
|
return None
|
|
|
|
def close(self) -> None:
|
|
self.close_calls += 1
|
|
|
|
|
|
class StreamingTurn:
|
|
sample_rate_hz = 24_000
|
|
retain_pcm = False
|
|
|
|
def __init__(self, events: list[str]) -> None:
|
|
self.events = events
|
|
self.closed = False
|
|
|
|
@property
|
|
def cleanup_complete(self) -> bool:
|
|
return self.closed
|
|
|
|
def enqueue_pcm(self, chunk: bytes) -> None:
|
|
self.events.append(f"audio-{chunk.decode('ascii')}")
|
|
|
|
def finish(self, capture: Capture | None = None) -> str:
|
|
self.events.append("finish")
|
|
self.assert_capture = capture
|
|
return "cloud result"
|
|
|
|
def poll_error(self) -> None:
|
|
return None
|
|
|
|
def cancel(self) -> None:
|
|
self.events.append("cancel")
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class StreamingRecorder(FakeRecorder):
|
|
def __init__(self, sink: StreamingTurn, events: list[str]) -> None:
|
|
super().__init__()
|
|
self.sink = sink
|
|
self.events = events
|
|
self.capture = Capture(b"", -15, 15, b"")
|
|
|
|
def start(self) -> None:
|
|
super().start()
|
|
self.sink.enqueue_pcm(b"a")
|
|
self.sink.enqueue_pcm(b"b")
|
|
|
|
def stop(self) -> Capture:
|
|
self.events.append("recorder-stop")
|
|
return super().stop()
|
|
|
|
|
|
class StreamingRecorderFactory:
|
|
def __init__(self, events: list[str]) -> None:
|
|
self.events = events
|
|
self.sample_rate_hz: int | None = None
|
|
self.instance: StreamingRecorder | None = None
|
|
|
|
def __call__(
|
|
self,
|
|
*,
|
|
sample_rate_hz: int,
|
|
pcm_sink: StreamingTurn,
|
|
retain_pcm: bool,
|
|
) -> StreamingRecorder:
|
|
self.sample_rate_hz = sample_rate_hz
|
|
if retain_pcm:
|
|
raise AssertionError("cloud capture must not retain PCM")
|
|
self.instance = StreamingRecorder(pcm_sink, self.events)
|
|
return self.instance
|
|
|
|
|
|
class FakeBackend:
|
|
sample_rate_hz = 16_000
|
|
descriptor = type("Descriptor", (), {"backend": "fake", "model": "fake-model"})()
|
|
|
|
def __init__(self, turn: DelayedCloseTurn | StreamingTurn) -> None:
|
|
self.turn = turn
|
|
|
|
def start_turn(self) -> DelayedCloseTurn:
|
|
return self.turn
|
|
|
|
def cancel_prepare(self) -> None:
|
|
return None
|
|
|
|
|
|
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": PROTOCOL_VERSION,
|
|
"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 AudioLevelTests(CoreTestCase):
|
|
def test_emits_owned_normalized_audio_levels_at_no_more_than_ten_hertz(self) -> None:
|
|
recorder = self.load_to_listening()
|
|
recorder.audio_level = 0.625
|
|
|
|
events: list[Outbound] = []
|
|
for _ in range(40):
|
|
self.clock.advance(0.025)
|
|
self.core.poll()
|
|
events.extend(self.core.drain_events())
|
|
|
|
self.assertGreaterEqual(len(events), 9)
|
|
self.assertLessEqual(len(events), 10)
|
|
self.assertEqual({event.client_id for event in events}, {"a"})
|
|
self.assertEqual({event.frame["type"] for event in events}, {"level"})
|
|
self.assertEqual({event.frame["level"] for event in events}, {0.625})
|
|
|
|
def test_clamps_invalid_levels_and_stops_emitting_after_listening(self) -> None:
|
|
recorder = self.load_to_listening()
|
|
for raw_level, expected in ((2.0, 1.0), (-1.0, 0.0), (float("nan"), 0.0)):
|
|
recorder.audio_level = raw_level
|
|
self.clock.advance(0.1)
|
|
self.core.poll()
|
|
level = self.frames(self.core.drain_events(), "a")
|
|
self.assertEqual(level[0]["level"], expected)
|
|
|
|
self.request("a", "stop", "stop", "g1")
|
|
self.clock.advance(0.2)
|
|
self.core.poll()
|
|
self.assertEqual(self.core.drain_events(), [])
|
|
|
|
|
|
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_realtime_backend_streams_ordered_pcm_and_finalizes_after_recorder_stop(self) -> None:
|
|
events: list[str] = []
|
|
turn = StreamingTurn(events)
|
|
recorders = StreamingRecorderFactory(events)
|
|
core = DictateServiceCore(
|
|
backend=FakeBackend(turn),
|
|
recorder_factory=recorders,
|
|
executor=self.executor,
|
|
clock=self.clock,
|
|
idle_seconds=900,
|
|
)
|
|
core.accepted("a")
|
|
core.handle_request("a", parse_request({"version": 2, "type": "hello", "requestId": "hello"}))
|
|
core.drain_events()
|
|
core.handle_request("a", parse_request({
|
|
"version": 2,
|
|
"type": "start",
|
|
"requestId": "start",
|
|
"generationId": "g1",
|
|
}))
|
|
self.executor.run_next()
|
|
core.poll()
|
|
self.assertEqual(recorders.sample_rate_hz, 24_000)
|
|
self.assertEqual(events, ["audio-a", "audio-b"])
|
|
core.drain_events()
|
|
|
|
core.handle_request("a", parse_request({
|
|
"version": 2,
|
|
"type": "stop",
|
|
"requestId": "stop",
|
|
"generationId": "g1",
|
|
}))
|
|
self.executor.run_next()
|
|
core.poll()
|
|
self.executor.run_next()
|
|
core.poll()
|
|
|
|
self.assertEqual(events, ["audio-a", "audio-b", "recorder-stop", "finish"])
|
|
result = [event.frame for event in core.drain_events() if event.frame["type"] == "result"]
|
|
self.assertEqual(result[0]["transcript"], "cloud result")
|
|
self.assertIsNone(core.active_generation_id)
|
|
|
|
def test_cancel_after_prepare_completion_tracks_turn_cleanup_before_release(self) -> None:
|
|
events: list[str] = []
|
|
turn = StreamingTurn(events)
|
|
core = DictateServiceCore(
|
|
backend=FakeBackend(turn),
|
|
recorder_factory=StreamingRecorderFactory(events),
|
|
executor=self.executor,
|
|
clock=self.clock,
|
|
idle_seconds=900,
|
|
)
|
|
core.accepted("a")
|
|
core.handle_request("a", parse_request({"version": 2, "type": "hello", "requestId": "hello"}))
|
|
core.drain_events()
|
|
core.handle_request("a", parse_request({
|
|
"version": 2,
|
|
"type": "start",
|
|
"requestId": "start",
|
|
"generationId": "g1",
|
|
}))
|
|
self.executor.run_next()
|
|
# The future has completed, but the socket loop has not yet applied it.
|
|
core.handle_request("a", parse_request({
|
|
"version": 2,
|
|
"type": "cancel",
|
|
"requestId": "cancel",
|
|
"generationId": "g1",
|
|
}))
|
|
|
|
core.poll()
|
|
|
|
self.assertIn("cancel", events)
|
|
self.assertTrue(turn.cleanup_complete)
|
|
self.assertIsNone(core.active_generation_id)
|
|
|
|
def test_completed_turn_retries_bounded_cleanup_before_releasing_generation(self) -> None:
|
|
turn = DelayedCloseTurn()
|
|
core = DictateServiceCore(
|
|
backend=FakeBackend(turn),
|
|
recorder_factory=self.recorders,
|
|
executor=self.executor,
|
|
clock=self.clock,
|
|
idle_seconds=900,
|
|
)
|
|
core.accepted("a")
|
|
core.handle_request("a", parse_request({"version": 2, "type": "hello", "requestId": "hello"}))
|
|
core.drain_events()
|
|
core.handle_request("a", parse_request({
|
|
"version": 2,
|
|
"type": "start",
|
|
"requestId": "start",
|
|
"generationId": "g1",
|
|
}))
|
|
self.executor.run_next()
|
|
core.poll()
|
|
core.drain_events()
|
|
core.handle_request("a", parse_request({
|
|
"version": 2,
|
|
"type": "stop",
|
|
"requestId": "stop",
|
|
"generationId": "g1",
|
|
}))
|
|
self.executor.run_next()
|
|
core.poll()
|
|
self.executor.run_next()
|
|
|
|
core.poll()
|
|
self.assertEqual(core.active_generation_id, "g1")
|
|
self.assertEqual(turn.close_calls, 2)
|
|
|
|
core.poll()
|
|
self.assertIsNone(core.active_generation_id)
|
|
self.assertEqual(turn.close_calls, 3)
|
|
|
|
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": PROTOCOL_VERSION, "type": "hello", "requestId": "h"}))
|
|
core.drain_events()
|
|
core.handle_request("a", parse_request({"version": PROTOCOL_VERSION, "type": "start", "requestId": "s", "generationId": "g"}))
|
|
core.drain_events()
|
|
core.handle_request("a", parse_request({"version": PROTOCOL_VERSION, "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()
|
|
self.executor.run_next()
|
|
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.executor.run_next()
|
|
self.core.poll()
|
|
failed.extend(self.frames(self.core.drain_events(), "a"))
|
|
|
|
self.assertEqual(failed[0]["code"], "recorder_exit")
|
|
self.assertEqual(failed[-1]["state"], "idle")
|
|
self.assertEqual(self.executor.pending, [])
|
|
|
|
|
|
def test_recorder_spawn_failure_releases_generation_without_cleanup_loop(self) -> None:
|
|
core = DictateServiceCore(
|
|
model=self.model,
|
|
recorder_factory=SpawnFailRecorderFactory(),
|
|
executor=self.executor,
|
|
clock=self.clock,
|
|
idle_seconds=900,
|
|
)
|
|
core.accepted("a")
|
|
core.handle_request("a", parse_request({
|
|
"version": 2, "type": "hello", "requestId": "hello"
|
|
}))
|
|
core.drain_events()
|
|
core.handle_request("a", parse_request({
|
|
"version": 2,
|
|
"type": "start",
|
|
"requestId": "start",
|
|
"generationId": "g1",
|
|
}))
|
|
self.executor.run_next()
|
|
|
|
core.poll()
|
|
|
|
events = self.frames(core.drain_events(), "a")
|
|
self.assertEqual(events[-2]["code"], "recorder_spawn")
|
|
self.assertEqual(events[-1]["state"], "idle")
|
|
self.assertIsNone(core.active_generation_id)
|
|
self.assertEqual(self.executor.pending, [])
|
|
|
|
def test_unreaped_recorder_keeps_generation_busy_until_cleanup_completes(self) -> None:
|
|
recorders = DelayedCleanupRecorderFactory()
|
|
core = DictateServiceCore(
|
|
model=self.model,
|
|
recorder_factory=recorders,
|
|
executor=self.executor,
|
|
clock=self.clock,
|
|
idle_seconds=900,
|
|
)
|
|
core.accepted("a")
|
|
core.handle_request("a", parse_request({
|
|
"version": 2, "type": "hello", "requestId": "hello"
|
|
}))
|
|
core.drain_events()
|
|
core.handle_request("a", parse_request({
|
|
"version": 2,
|
|
"type": "start",
|
|
"requestId": "start",
|
|
"generationId": "g1",
|
|
}))
|
|
self.executor.run_next()
|
|
core.poll()
|
|
core.drain_events()
|
|
core.handle_request("a", parse_request({
|
|
"version": 2,
|
|
"type": "stop",
|
|
"requestId": "stop",
|
|
"generationId": "g1",
|
|
}))
|
|
self.executor.run_next()
|
|
|
|
core.poll()
|
|
self.assertEqual(core.active_generation_id, "g1")
|
|
# First reap the cancellation worker scheduled by the terminal stop.
|
|
self.executor.run_next()
|
|
core.poll()
|
|
for attempt in range(1, 4):
|
|
self.executor.run_next()
|
|
core.poll()
|
|
if attempt < 3:
|
|
self.assertEqual(core.active_generation_id, "g1")
|
|
self.assertIsNone(core.active_generation_id)
|
|
self.assertEqual(recorders.instance.cleanup_attempts, 3)
|
|
|
|
|
|
class SocketLoopRegressionTests(unittest.TestCase):
|
|
def test_output_overflow_does_not_use_the_closed_socket_from_same_selector_event(self) -> None:
|
|
frame = {"padding": "x" * (MAX_FRAME_BYTES - len(encode_frame({"padding": ""})))}
|
|
self.assertEqual(len(encode_frame(frame)), MAX_FRAME_BYTES)
|
|
client = _FakeClientSocket()
|
|
listener = _FakeListener(client)
|
|
selector = _FakeSelector(client)
|
|
core = _FakeLoopCore(
|
|
Outbound("client-1", frame),
|
|
)
|
|
|
|
with patch("pi_dictate_service.server.selectors.DefaultSelector", return_value=selector):
|
|
serve(listener, core)
|
|
|
|
self.assertEqual(client.send_calls, 0)
|
|
self.assertTrue(client.closed)
|
|
self.assertTrue(listener.closed)
|
|
|
|
def test_close_does_not_wait_for_owned_executor_workers(self) -> None:
|
|
executor = _ShutdownRecordingExecutor()
|
|
with patch("pi_dictate_service.server._DaemonExecutor", return_value=executor):
|
|
core = DictateServiceCore(model=FakeModel())
|
|
core.accepted("a")
|
|
core.handle_request(
|
|
"a", parse_request({"version": PROTOCOL_VERSION, "type": "hello", "requestId": "h"})
|
|
)
|
|
core.drain_events()
|
|
core.handle_request(
|
|
"a",
|
|
parse_request(
|
|
{
|
|
"version": PROTOCOL_VERSION,
|
|
"type": "start",
|
|
"requestId": "s",
|
|
"generationId": "g",
|
|
}
|
|
),
|
|
)
|
|
core.close()
|
|
|
|
self.assertEqual(executor.shutdown_args, (False, True))
|
|
|
|
def test_real_owned_executor_does_not_block_interpreter_exit(self) -> None:
|
|
program = """
|
|
import threading
|
|
import time
|
|
from pi_dictate_service.protocol import parse_request
|
|
from pi_dictate_service.server import DictateServiceCore
|
|
|
|
class StuckModel:
|
|
loaded = False
|
|
def load(self):
|
|
threading.Event().wait()
|
|
def transcribe(self, _samples):
|
|
return ''
|
|
|
|
core = DictateServiceCore(model=StuckModel())
|
|
core.accepted('a')
|
|
core.handle_request('a', parse_request({'version': 2, 'type': 'hello', 'requestId': 'h'}))
|
|
core.drain_events()
|
|
core.handle_request('a', parse_request({
|
|
'version': 2, 'type': 'start', 'requestId': 's', 'generationId': 'g'
|
|
}))
|
|
time.sleep(0.05)
|
|
core.close()
|
|
print('closed')
|
|
"""
|
|
|
|
result = subprocess.run(
|
|
[sys.executable, "-c", program],
|
|
text=True,
|
|
capture_output=True,
|
|
timeout=2,
|
|
check=False,
|
|
)
|
|
|
|
self.assertEqual(result.returncode, 0, result.stderr)
|
|
self.assertEqual(result.stdout.strip(), "closed")
|
|
|
|
|
|
@dataclass
|
|
class _FakeSelectorKey:
|
|
data: str | None
|
|
|
|
|
|
class _FakeSelector:
|
|
def __init__(self, client: "_FakeClientSocket") -> None:
|
|
self._events = [
|
|
[(_FakeSelectorKey(None), selectors.EVENT_READ)],
|
|
[(_FakeSelectorKey("client-1"), selectors.EVENT_READ | selectors.EVENT_WRITE)],
|
|
]
|
|
self.client = client
|
|
self.unregistered: list[object] = []
|
|
self.closed = False
|
|
|
|
def register(self, fileobj: object, _events: int, _data: object = None) -> None:
|
|
return None
|
|
|
|
def unregister(self, fileobj: object) -> None:
|
|
self.unregistered.append(fileobj)
|
|
|
|
def modify(self, fileobj: object, _events: int, _data: object = None) -> None:
|
|
if fileobj is self.client and self.client.closed:
|
|
raise OSError(errno.EBADF, "bad file descriptor")
|
|
|
|
def select(self, timeout: float | None = None) -> list[tuple[_FakeSelectorKey, int]]:
|
|
return self._events.pop(0)
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class _FakeClientSocket:
|
|
def __init__(self) -> None:
|
|
self.closed = False
|
|
self.send_calls = 0
|
|
self._read = True
|
|
|
|
def setblocking(self, _value: bool) -> None:
|
|
return None
|
|
|
|
def recv(self, _size: int) -> bytes:
|
|
if self._read:
|
|
self._read = False
|
|
return b"{broken}\n"
|
|
return b""
|
|
|
|
def send(self, _data: bytearray) -> int:
|
|
self.send_calls += 1
|
|
raise OSError(errno.EBADF, "bad file descriptor")
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class _FakeListener:
|
|
def __init__(self, client: _FakeClientSocket) -> None:
|
|
self.client = client
|
|
self.closed = False
|
|
self._accepted = False
|
|
|
|
def setblocking(self, _value: bool) -> None:
|
|
return None
|
|
|
|
def accept(self) -> tuple[_FakeClientSocket, object]:
|
|
if self._accepted:
|
|
raise BlockingIOError
|
|
self._accepted = True
|
|
return self.client, None
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class _FakeLoopCore:
|
|
def __init__(self, outbound: Outbound) -> None:
|
|
self.outbound = outbound
|
|
self.accepted_client: str | None = None
|
|
self._outbound_sent = False
|
|
self._exit_checks = 0
|
|
|
|
def should_exit(self) -> bool:
|
|
self._exit_checks += 1
|
|
return self._exit_checks >= 3
|
|
|
|
def poll(self) -> None:
|
|
return None
|
|
|
|
def drain_events(self) -> list[Outbound]:
|
|
if self.accepted_client is not None and not self._outbound_sent:
|
|
self._outbound_sent = True
|
|
return [self.outbound]
|
|
return []
|
|
|
|
def accepted(self, client_id: str) -> None:
|
|
self.accepted_client = client_id
|
|
|
|
def disconnect(self, _client_id: str) -> None:
|
|
return None
|
|
|
|
def close(self) -> None:
|
|
return None
|
|
|
|
|
|
class _ShutdownRecordingExecutor:
|
|
def __init__(self) -> None:
|
|
self.shutdown_args: tuple[bool, bool] | None = None
|
|
|
|
def submit(self, _function: object, *_args: object) -> Future[object]:
|
|
return Future()
|
|
|
|
def shutdown(self, *, wait: bool = True, cancel_futures: bool = False) -> None:
|
|
self.shutdown_args = (wait, cancel_futures)
|
|
if wait:
|
|
raise AssertionError("shutdown waited for a hung worker")
|
|
|
|
|
|
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()
|