334 lines
12 KiB
Python
334 lines
12 KiB
Python
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()
|