jb/harness-tools
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(dictate): add voice-reactive listening waveform

Browse Source
This commit is contained in:
jb
2026-08-26 13:40:47 -04:00
parent e234468a73
commit d1a811a5c5
16 changed files with 383 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
services/dictate/src/pi_dictate_service/model.py
+4 -3
View File
@@ -86,12 +86,13 @@ class NemoParakeetModel:
os.environ["NEMO_OFFLINE"] = "1"
import numpy
from nemo.collections.asr.models import ASRModel
import torch
from nemo.collections.asr.models import EncDecRNNTBPEModel
from omegaconf import open_dict
model = ASRModel.restore_from(
model = EncDecRNNTBPEModel.restore_from(
restore_path=str(self._model_path),
map_location="cuda",
map_location=torch.device("cuda"),
)
model.to("cuda")
model.eval()
services/dictate/src/pi_dictate_service/protocol.py
+2 -2
View File
@@ -1,4 +1,4 @@
"""Bounded newline-delimited JSON protocol primitives for dictate protocol v1."""
"""Bounded newline-delimited JSON protocol primitives for dictate protocol v2."""
from __future__ import annotations
@@ -6,7 +6,7 @@ from dataclasses import dataclass
import json
from typing import Any, Mapping
PROTOCOL_VERSION = 1
PROTOCOL_VERSION = 2
MAX_FRAME_BYTES = 64 * 1024
MAX_TRANSCRIPT_BYTES = 32 * 1024
MAX_ERROR_DETAIL_BYTES = 2 * 1024
services/dictate/src/pi_dictate_service/recorder.py
+29 -1
View File
@@ -2,10 +2,13 @@
from __future__ import annotations
from array import array
from dataclasses import dataclass
import math
import os
import signal
import subprocess
import sys
import threading
import time
from typing import BinaryIO, Callable, Protocol
@@ -24,6 +27,7 @@ PW_RECORD_COMMAND = [
"-",
]
MAX_STDERR_BYTES = 64 * 1024
AUDIO_LEVEL_WINDOW_BYTES = 3_200 # 100 ms of mono 16 kHz signed-16 PCM
class Process(Protocol):
@@ -95,12 +99,18 @@ class PwRecorder:
self._duration_limited = False
self._duration_stop_at: float | None = None
self._reported = False
self._audio_level = 0.0
@property
def buffered_bytes(self) -> int:
with self._lock:
return len(self._pcm)
@property
def audio_level(self) -> float:
with self._lock:
return self._audio_level
def start(self) -> None:
if self._process is not None:
raise RecorderError("invalid_state", "recorder is already started")
@@ -223,10 +233,15 @@ class PwRecorder:
def _drain_stdout(self) -> None:
process = self._require_process()
level_window = bytearray()
while True:
chunk = process.stdout.read(64 * 1024)
chunk = process.stdout.read(AUDIO_LEVEL_WINDOW_BYTES)
if not chunk:
return
level_window.extend(chunk)
while len(level_window) >= AUDIO_LEVEL_WINDOW_BYTES:
self._update_audio_level(level_window[:AUDIO_LEVEL_WINDOW_BYTES])
del level_window[:AUDIO_LEVEL_WINDOW_BYTES]
with self._lock:
if self._duration_limited:
continue
@@ -240,6 +255,19 @@ class PwRecorder:
if process.poll() is None:
self._signal_group(process, signal.SIGTERM)
def _update_audio_level(self, chunk: bytes) -> None:
sample_bytes = len(chunk) - (len(chunk) % 2)
if sample_bytes == 0:
return
samples = array("h")
samples.frombytes(chunk[:sample_bytes])
if sys.byteorder != "little":
samples.byteswap()
rms = math.sqrt(sum(sample * sample for sample in samples) / len(samples)) / 32768.0
level = 0.0 if rms == 0 else max(0.0, min(1.0, (20 * math.log10(rms) + 60) / 60))
with self._lock:
self._audio_level = max(level, self._audio_level * 0.65)
def _drain_stderr(self) -> None:
process = self._require_process()
while True:
services/dictate/src/pi_dictate_service/server.py
+32 -1
View File
@@ -5,6 +5,7 @@ from __future__ import annotations
from collections.abc import Sequence
from concurrent.futures import Executor, Future, ThreadPoolExecutor
from dataclasses import dataclass
import math
import os
from queue import Empty, SimpleQueue
import selectors
@@ -30,9 +31,13 @@ from .recorder import Capture, PwRecorder, RecorderError
DEFAULT_IDLE_SECONDS = 15 * 60
MAX_CLIENT_CONNECTIONS = 16
MAX_CONNECTION_OUTPUT_BYTES = MAX_FRAME_BYTES
LEVEL_EVENT_INTERVAL_SECONDS = 0.1
class Recorder(Protocol):
@property
def audio_level(self) -> float: ...
def start(self) -> None: ...
def stop(self) -> Capture: ...
def cancel(self) -> None: ...
@@ -53,6 +58,7 @@ class _Generation:
state: str
recorder: Recorder | None = None
tombstoned: bool = False
last_level_event_at: float | None = None
@dataclass(frozen=True)
@@ -156,6 +162,8 @@ class DictateServiceCore:
if failure is not None:
self._async_error(active, failure.code, failure.detail)
self._finish_generation(active, notify_idle=True)
return
self._emit_audio_level(active)
def drain_events(self) -> list[Outbound]:
events, self._events = self._events, []
@@ -216,8 +224,31 @@ class DictateServiceCore:
self._finish_generation(active, notify_idle=True)
return
active.recorder = recorder
active.last_level_event_at = self._clock()
self._state(active, "listening")
def _emit_audio_level(self, active: _Generation) -> None:
now = self._clock()
if (
active.last_level_event_at is not None
and now - active.last_level_event_at < LEVEL_EVENT_INTERVAL_SECONDS
):
return
level = active.recorder.audio_level if active.recorder is not None else 0.0
if not math.isfinite(level):
level = 0.0
active.last_level_event_at = now
self._emit(
active.owner,
{
"version": PROTOCOL_VERSION,
"type": "level",
"requestId": active.request_id,
"generationId": active.generation_id,
"level": round(max(0.0, min(1.0, level)), 3),
},
)
def _stop(self, client_id: str, request: Request) -> None:
active = self._owned_generation(client_id, request)
if active is None:
@@ -433,7 +464,7 @@ def inherited_listener(environ: dict[str, str] | None = None) -> socket.socket:
def serve(listener: socket.socket, core: DictateServiceCore) -> None:
"""Serve protocol v1 on an explicit listener until the core becomes idle."""
"""Serve protocol v2 on an explicit listener until the core becomes idle."""
selector = selectors.DefaultSelector()
connections: dict[str, _Connection] = {}
sequence = 0
services/dictate/tests/test_model.py
+17 -11
View File
@@ -35,13 +35,15 @@ class FakeNemoModel:
return [type("Hypothesis", (), {"text": "hello locally"})()]
class FakeASRModel:
class FakeConcreteParakeetModel:
restored_path: str | None = None
restored_location: object | None = None
restored: FakeNemoModel | None = None
@classmethod
def restore_from(cls, restore_path: str, map_location: str) -> FakeNemoModel:
def restore_from(cls, restore_path: str, map_location: object) -> FakeNemoModel:
cls.restored_path = restore_path
cls.restored_location = map_location
self = FakeNemoModel()
cls.restored = self
return self
@@ -69,28 +71,29 @@ class NemoParakeetModelTests(unittest.TestCase):
boost_alpha=0.5,
expected_sha256=hashlib.sha256(b"model").hexdigest(),
)
self.assertIsNone(FakeASRModel.restored_path)
self.assertIsNone(FakeConcreteParakeetModel.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(FakeConcreteParakeetModel.restored_path, str(model_path))
self.assertEqual(FakeConcreteParakeetModel.restored_location, "device:cuda")
self.assertTrue(FakeConcreteParakeetModel.restored.eval_called)
self.assertEqual(FakeConcreteParakeetModel.restored.to_device, "cuda")
self.assertEqual(transcript, "hello locally")
self.assertEqual(FakeASRModel.restored.audio, [([0.0, 0.5], "float32")])
self.assertEqual(FakeConcreteParakeetModel.restored.audio, [([0.0, 0.5], "float32")])
self.assertEqual(
FakeASRModel.restored.decoding_config["greedy"]["boosting_tree"]["key_phrases_list"],
FakeConcreteParakeetModel.restored.decoding_config["greedy"]["boosting_tree"]["key_phrases_list"],
["JSON", "npm test"],
)
self.assertEqual(
FakeASRModel.restored.decoding_config["greedy"]["boosting_tree_alpha"],
FakeConcreteParakeetModel.restored.decoding_config["greedy"]["boosting_tree_alpha"],
0.5,
)
self.assertEqual(
FakeASRModel.restored.transcribe_options,
FakeConcreteParakeetModel.restored.transcribe_options,
{"batch_size": 1, "num_workers": 0, "verbose": False},
)
@@ -123,12 +126,14 @@ class NemoParakeetModelTests(unittest.TestCase):
collections = ModuleType("nemo.collections")
asr = ModuleType("nemo.collections.asr")
models = ModuleType("nemo.collections.asr.models")
models.ASRModel = FakeASRModel
models.EncDecRNNTBPEModel = FakeConcreteParakeetModel
asr.models = models
collections.asr = asr
nemo.collections = collections
omegaconf = ModuleType("omegaconf")
omegaconf.open_dict = lambda _config: nullcontext()
torch = ModuleType("torch")
torch.device = lambda name: f"device:{name}"
return {
"nemo": nemo,
"nemo.collections": collections,
@@ -136,6 +141,7 @@ class NemoParakeetModelTests(unittest.TestCase):
"nemo.collections.asr.models": models,
"numpy": FakeNumpy("numpy"),
"omegaconf": omegaconf,
"torch": torch,
}
services/dictate/tests/test_protocol_audio.py
+13 -12
View File
@@ -7,6 +7,7 @@ from pi_dictate_service.protocol import (
MAX_FRAME_BYTES,
MAX_ID_BYTES,
MAX_TRANSCRIPT_BYTES,
PROTOCOL_VERSION,
FrameDecoder,
ProtocolError,
encode_frame,
@@ -18,10 +19,10 @@ class FrameDecoderTests(unittest.TestCase):
def test_decodes_fragmented_and_multiple_frames(self) -> None:
decoder = FrameDecoder()
self.assertEqual(decoder.feed(b'{"version":1'), [])
self.assertEqual(decoder.feed(f'{{"version":{PROTOCOL_VERSION}'.encode()), [])
frames = decoder.feed(
b',"type":"ping","requestId":"one"}\n'
b'{"version":1,"type":"ping","requestId":"two"}\n'
+ f'{{"version":{PROTOCOL_VERSION},"type":"ping","requestId":"two"}}\n'.encode()
)
self.assertEqual([frame["requestId"] for frame in frames], ["one", "two"])
@@ -41,14 +42,14 @@ class FrameDecoderTests(unittest.TestCase):
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"}
{"version": PROTOCOL_VERSION, "type": "start", "requestId": "r1", "generationId": "g1"}
)
self.assertEqual(request.generation_id, "g1")
with self.assertRaises(ProtocolError) as unknown:
parse_request(
{
"version": 1,
"version": PROTOCOL_VERSION,
"type": "start",
"requestId": "r1",
"generationId": "g1",
@@ -61,7 +62,7 @@ class RequestValidationTests(unittest.TestCase):
with self.assertRaises(ProtocolError) as long_id:
parse_request(
{
"version": 1,
"version": PROTOCOL_VERSION,
"type": "ping",
"requestId": "é" * ((MAX_ID_BYTES // 2) + 1),
}
@@ -69,7 +70,7 @@ class RequestValidationTests(unittest.TestCase):
self.assertEqual(long_id.exception.code, "invalid_request")
with self.assertRaises(ProtocolError) as unsupported:
parse_request({"version": 2, "type": "hello", "requestId": "r2"})
parse_request({"version": PROTOCOL_VERSION + 1, "type": "hello", "requestId": "r2"})
self.assertEqual(unsupported.exception.code, "unsupported_version")
self.assertTrue(unsupported.exception.close_connection)
@@ -78,12 +79,12 @@ class RequestValidationTests(unittest.TestCase):
with self.subTest(request_type=request_type):
with self.assertRaises(ProtocolError) as failed:
parse_request(
{"version": 1, "type": request_type, "requestId": "request"}
{"version": PROTOCOL_VERSION, "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"})
parse_request({"version": PROTOCOL_VERSION, "type": "ping", "requestId": "\ud800"})
self.assertEqual(surrogate.exception.request_id, "")
huge_integer = b'{"version":' + (b"9" * 5000) + b',"type":"hello","requestId":"r"}\n'
@@ -96,13 +97,13 @@ class RequestValidationTests(unittest.TestCase):
oversized = "x" * (MAX_ID_BYTES + 1)
with self.assertRaises(ProtocolError) as unsupported:
parse_request(
{"version": 2, "type": "hello", "requestId": oversized}
{"version": PROTOCOL_VERSION + 1, "type": "hello", "requestId": oversized}
)
self.assertEqual(unsupported.exception.request_id, "")
encoded = encode_frame(
{
"version": 1,
"version": PROTOCOL_VERSION,
"type": "error",
"requestId": oversized,
"generationId": "\ud800",
@@ -118,7 +119,7 @@ class RequestValidationTests(unittest.TestCase):
with self.assertRaises(ProtocolError) as transcript:
encode_frame(
{
"version": 1,
"version": PROTOCOL_VERSION,
"type": "result",
"generationId": "g1",
"transcript": "é" * ((MAX_TRANSCRIPT_BYTES // 2) + 1),
@@ -128,7 +129,7 @@ class RequestValidationTests(unittest.TestCase):
encoded = encode_frame(
{
"version": 1,
"version": PROTOCOL_VERSION,
"type": "error",
"requestId": "r1",
"code": "failed",
services/dictate/tests/test_recorder.py
+69 -3
View File
@@ -1,19 +1,34 @@
import io
import math
import subprocess
import threading
import unittest
from pi_dictate_service.recorder import PW_RECORD_COMMAND, PwRecorder, RecorderError
from pi_dictate_service.recorder import (
AUDIO_LEVEL_WINDOW_BYTES,
PW_RECORD_COMMAND,
PwRecorder,
RecorderError,
)
class TrackingStream(io.BytesIO):
def __init__(self, value: bytes, events: list[str], name: str) -> None:
def __init__(
self,
value: bytes,
events: list[str],
name: str,
max_read_size: int | None = None,
) -> None:
super().__init__(value)
self._events = events
self._name = name
self._max_read_size = max_read_size
self.eof = threading.Event()
def read(self, size: int = -1) -> bytes:
if self._max_read_size is not None:
size = min(size, self._max_read_size)
value = super().read(size)
if not value and not self.eof.is_set():
self._events.append(f"{self._name}-eof")
@@ -91,10 +106,11 @@ class PwRecorderTests(unittest.TestCase):
stderr_limit: int = 64 * 1024,
clock: Clock | None = None,
terminate_returncode: int = -15,
stdout_read_size: int | None = None,
) -> tuple[PwRecorder, FakeProcess, ProcessFactory, list[str]]:
events: list[str] = []
process = FakeProcess(
TrackingStream(pcm, events, "stdout"),
TrackingStream(pcm, events, "stdout", stdout_read_size),
TrackingStream(stderr, events, "stderr"),
returncode=returncode,
timeout=timeout,
@@ -138,6 +154,56 @@ class PwRecorderTests(unittest.TestCase):
self.assertIn("stdout-eof", events)
self.assertIn("wait", events)
def test_tracks_normalized_audio_level_from_pcm_without_exposing_audio(self) -> None:
recorder, process, _factory, _events = self.make_recorder(
pcm=b"\xff\x7f" * (AUDIO_LEVEL_WINDOW_BYTES // 2),
max_bytes=AUDIO_LEVEL_WINDOW_BYTES,
)
recorder.start()
self.assertTrue(process.stdout.eof.wait(1))
self.assertGreater(recorder.audio_level, 0.99)
self.assertLessEqual(recorder.audio_level, 1.0)
recorder.cancel()
def test_audio_level_accumulates_fragmented_reads_into_100ms_windows(self) -> None:
loud_sample = (32767).to_bytes(2, "little", signed=True)
pcm = (
loud_sample * (AUDIO_LEVEL_WINDOW_BYTES // 4)
+ b"\x00" * (AUDIO_LEVEL_WINDOW_BYTES // 2)
+ loud_sample * 400
)
recorder, process, _factory, _events = self.make_recorder(
pcm=pcm,
max_bytes=len(pcm),
stdout_read_size=800,
)
recorder.start()
self.assertTrue(process.stdout.eof.wait(1))
expected_rms = (0.5**0.5) * 32767 / 32768
expected_level = (20 * math.log10(expected_rms) + 60) / 60
self.assertAlmostEqual(recorder.audio_level, expected_level, places=3)
recorder.cancel()
def test_audio_level_maps_silence_known_rms_decay_and_partial_samples(self) -> None:
recorder = PwRecorder()
recorder._update_audio_level(b"\x00\x00" * 16)
self.assertEqual(recorder.audio_level, 0.0)
amplitude = 3277 # approximately -20 dBFS
recorder._update_audio_level(amplitude.to_bytes(2, "little", signed=True) * 16)
self.assertAlmostEqual(recorder.audio_level, 2 / 3, places=3)
recorder._update_audio_level(b"\x00\x00" * 16)
self.assertAlmostEqual(recorder.audio_level, (2 / 3) * 0.65, places=3)
previous = recorder.audio_level
recorder._update_audio_level(b"\xff")
self.assertEqual(recorder.audio_level, previous)
def test_accepts_pw_record_code_one_after_intentional_sigterm(self) -> None:
recorder, process, _factory, _events = self.make_recorder(terminate_returncode=1)
recorder.start()
services/dictate/tests/test_server_core.py
+38 -5
View File
@@ -3,7 +3,7 @@ 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.protocol import MAX_TRANSCRIPT_BYTES, PROTOCOL_VERSION, parse_request
from pi_dictate_service.recorder import Capture, RecorderError
from pi_dictate_service.server import DictateServiceCore, Outbound
@@ -73,6 +73,7 @@ class FakeRecorder:
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
@@ -124,7 +125,7 @@ class CoreTestCase(unittest.TestCase):
generation_id: str | None = None,
) -> list[dict[str, object]]:
frame: dict[str, object] = {
"version": 1,
"version": PROTOCOL_VERSION,
"type": request_type,
"requestId": request_id,
}
@@ -217,6 +218,38 @@ class CancellationTests(CoreTestCase):
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()
@@ -271,11 +304,11 @@ class StopAndInferenceTests(CoreTestCase):
executor=self.executor,
clock=self.clock,
)
core.handle_request("a", parse_request({"version": 1, "type": "hello", "requestId": "h"}))
core.handle_request("a", parse_request({"version": PROTOCOL_VERSION, "type": "hello", "requestId": "h"}))
core.drain_events()
core.handle_request("a", parse_request({"version": 1, "type": "start", "requestId": "s", "generationId": "g"}))
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": 1, "type": "stop", "requestId": "t", "generationId": "g"}))
core.handle_request("a", parse_request({"version": PROTOCOL_VERSION, "type": "stop", "requestId": "t", "generationId": "g"}))
self.executor.run_next()
core.poll()
core.drain_events()
services/dictate/tests/test_unix_server.py
+14 -14
View File
@@ -73,24 +73,24 @@ class UnixLoopTests(unittest.TestCase):
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'{"version":2,"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')
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":1,"type":"hello","requestId":"hello"}\n')
client.sendall(b'{"version":2,"type":"hello","requestId":"hello"}\n')
self.receive_frame(client)
client.sendall(
b'{"version":1,"type":"stop","requestId":"stop","generationId":"g"}\n'
b'{"version":2,"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')
client.sendall(b'{"version":2,"type":"ping","requestId":"ping"}\n')
self.assertEqual(self.receive_frame(client)["type"], "pong")
client.sendall(b'{broken}\n')
@@ -101,9 +101,9 @@ class UnixLoopTests(unittest.TestCase):
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')
client.sendall(b'{"version":2,"type":[],"requestId":"bad-type"}\n')
self.assertEqual(self.receive_frame(client)["code"], "invalid_request")
client.sendall(b'{"version":1,"type":"ping","requestId":"\\ud800"}\n')
client.sendall(b'{"version":2,"type":"ping","requestId":"\\ud800"}\n')
self.assertEqual(self.receive_frame(client)["code"], "invalid_request")
with self.connect() as client:
@@ -115,14 +115,14 @@ class UnixLoopTests(unittest.TestCase):
self.assertEqual(client.recv(1), b"")
with self.connect() as client:
client.sendall(b'{"version":1,"type":"hello","requestId":"healthy"}\n')
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":1,"type":"hello","requestId":"hello"}\n')
client.sendall(b'{"version":2,"type":"hello","requestId":"hello"}\n')
self.receive_frame(client)
ping = b'{"version":1,"type":"ping","requestId":"p"}\n'
ping = b'{"version":2,"type":"ping","requestId":"p"}\n'
try:
client.sendall(ping * 1200)
except BrokenPipeError:
@@ -139,7 +139,7 @@ class UnixLoopTests(unittest.TestCase):
self.assertLessEqual(len(received), 64 * 1024)
with self.connect() as client:
client.sendall(b'{"version":1,"type":"hello","requestId":"healthy"}\n')
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:
@@ -148,7 +148,7 @@ class UnixLoopTests(unittest.TestCase):
time.sleep(0.1)
overflow = self.connect()
try:
overflow.sendall(b'{"version":1,"type":"hello","requestId":"overflow"}\n')
overflow.sendall(b'{"version":2,"type":"hello","requestId":"overflow"}\n')
with self.assertRaises((ConnectionResetError, BrokenPipeError)):
if overflow.recv(1) == b"":
raise ConnectionResetError
@@ -161,12 +161,12 @@ class UnixLoopTests(unittest.TestCase):
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.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":1,"type":"hello","requestId":"hello"}\n')
client.sendall(b'{"version":2,"type":"hello","requestId":"hello"}\n')
self.assertEqual(self.receive_frame(client)["type"], "pong")