harness-tools/services/dictate/tests/test_protocol_audio.py

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,
    PROTOCOL_VERSION,
    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(f'{{"version":{PROTOCOL_VERSION}'.encode()), [])
        frames = decoder.feed(
            b',"type":"ping","requestId":"one"}\n'
            + f'{{"version":{PROTOCOL_VERSION},"type":"ping","requestId":"two"}}\n'.encode()
        )

        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": PROTOCOL_VERSION, "type": "start", "requestId": "r1", "generationId": "g1"}
        )
        self.assertEqual(request.generation_id, "g1")

        with self.assertRaises(ProtocolError) as unknown:
            parse_request(
                {
                    "version": PROTOCOL_VERSION,
                    "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": PROTOCOL_VERSION,
                    "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": PROTOCOL_VERSION + 1, "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": PROTOCOL_VERSION, "type": request_type, "requestId": "request"}
                    )
                self.assertEqual(failed.exception.code, "invalid_request")

        with self.assertRaises(ProtocolError) as surrogate:
            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'
        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": PROTOCOL_VERSION + 1, "type": "hello", "requestId": oversized}
            )
        self.assertEqual(unsupported.exception.request_id, "")

        encoded = encode_frame(
            {
                "version": PROTOCOL_VERSION,
                "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": PROTOCOL_VERSION,
                    "type": "result",
                    "generationId": "g1",
                    "transcript": "é" * ((MAX_TRANSCRIPT_BYTES // 2) + 1),
                }
            )
        self.assertEqual(transcript.exception.code, "transcript_too_large")

        encoded = encode_frame(
            {
                "version": PROTOCOL_VERSION,
                "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()