old-books/_pipeline/tests/test_pipeline.py

"""Tests. Everything audio here is a synthetic TEST signal made locally, never pilot audio.
Run: python3 -m unittest discover -s _pipeline/tests -v   (from the repo root)"""
import http.server
import json
import os
import sys
import tempfile
import threading
import unittest

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))

import acx_check
import assemble
import audiomeasure as am
import script
import settings
import voice

CFG = settings.load()


def make_test_piece(path, hz, seconds, level):
    """TEST signal: speech-like (harmonics, 4 Hz syllable rhythm, a little noise) at amplitude `level`."""
    expr = "%g*(0.6*sin(2*PI*%g*t)+0.3*sin(4*PI*%g*t)+0.1*sin(6*PI*%g*t))*(0.55+0.45*sin(2*PI*4*t))" % (level, hz, hz, hz)
    am.ffmpeg(["-f", "lavfi", "-i", "aevalsrc='%s':s=24000:d=%g" % (expr, seconds), "-ar", "24000", "-ac", "1", path])


class AudioTests(unittest.TestCase):
    def test_assemble_then_check_passes(self):
        with tempfile.TemporaryDirectory() as d:
            items = []
            for i, (hz, lvl) in enumerate([(120, 0.02), (150, 0.05), (130, 0.03)]):
                p = os.path.join(d, "t%d.wav" % i)
                make_test_piece(p, hz, 4, lvl)
                items.append({"wav": p, "gap_after_s": 0.7})
                if i == 0:
                    items.append({"cue": "chime", "gap_after_s": 0.25})
            out = os.path.join(d, "audio", "chapter-99.mp3")
            gain = assemble.assemble_items(CFG, items, out)
            m = acx_check.measure(CFG, out)
            print("\nTEST assembled: gain %.1f dB, measured %s" % (gain, {k: (round(v, 2) if isinstance(v, float) else v) for k, v in m.items()}))
            for name, value, passed in acx_check.judge(CFG, out, m):
                self.assertTrue(passed, "%s %s" % (name, value))

    def test_checker_fails_bad_files(self):
        with tempfile.TemporaryDirectory() as d:
            raw = os.path.join(d, "raw.wav")
            make_test_piece(raw, 130, 4, 1.4)                       # hot test signal, clips at full scale
            bad = os.path.join(d, "bad.mp3")                   # stereo, 128k, 22.05 kHz, no room tone
            am.ffmpeg(["-i", raw, "-ac", "2", "-ar", "22050", "-c:a", "libmp3lame", "-b:a", "128k", bad])
            failed = {n for n, v, p in acx_check.judge(CFG, bad, acx_check.measure(CFG, bad)) if not p}
            print("\nTEST bad file 1 fails on:", sorted(failed))
            self.assertTrue({"format", "rms dB", "noise floor dB", "head s", "tail s"} <= failed)
            hot = os.path.join(d, "hot.mp3")                   # correct format, but too loud
            am.ffmpeg(["-i", raw, "-ac", "1", "-ar", "44100", "-c:a", "libmp3lame", "-b:a", "192k", hot])
            m = acx_check.measure(CFG, hot)
            failed = {n for n, v, p in acx_check.judge(CFG, hot, m) if not p}
            print("TEST bad file 2 (peak %.1f dB) fails on:" % m["peak_db"], sorted(failed))
            self.assertIn("peak dB", failed)
            self.assertNotIn("format", failed)


class StubHandler(http.server.BaseHTTPRequestHandler):
    seen = []

    def do_POST(self):
        body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
        StubHandler.seen.append((self.path, body))
        if self.path.endswith("/audio/speech") or self.path == "/tts":
            with tempfile.TemporaryDirectory() as d:
                p = os.path.join(d, "x.wav")
                make_test_piece(p, 130, 1, 0.1)
                data = open(p, "rb").read()
            self.send_response(200)
            self.send_header("Content-Type", "audio/wav")
        else:
            data = json.dumps({"choices": [{"message": {"content": "A test summary."}}]}).encode()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def log_message(self, *a):
        pass


class ClientTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.server = http.server.HTTPServer(("127.0.0.1", 0), StubHandler)
        threading.Thread(target=cls.server.serve_forever, daemon=True).start()
        cls.url = "http://127.0.0.1:%d" % cls.server.server_port

    @classmethod
    def tearDownClass(cls):
        cls.server.shutdown()

    def cfg(self):
        c = json.loads(json.dumps(CFG))
        c["voice"]["endpoint"] = self.url + "/v1/audio/speech"
        c["voice"]["api"] = "openai"
        c["voice"]["extra"] = {}
        c["llm"].update(endpoint=self.url + "/v1/chat/completions", model="m")
        return c

    def test_voice_request_and_wav(self):
        StubHandler.seen.clear()
        data = voice.speak(self.cfg(), "Hello there.")
        path, body = StubHandler.seen[0]
        self.assertEqual(path, "/v1/audio/speech")
        self.assertEqual(body["input"], "Hello there.")
        self.assertEqual(body["response_format"], "wav")
        self.assertEqual(data[:4], b"RIFF")

    def test_locked_voice_uses_native_route_with_locked_settings(self):
        StubHandler.seen.clear()
        c = json.loads(json.dumps(CFG))
        c["voice"]["endpoint"] = self.url + "/tts"
        voice.speak(c, "Hello there.")
        path, body = StubHandler.seen[0]
        self.assertEqual(path, "/tts")
        self.assertEqual(body["text"], "Hello there.")
        self.assertEqual(body["predefined_voice_id"], "Michael.wav")
        self.assertEqual((body["exaggeration"], body["cfg_weight"], body["temperature"], body["seed"], body["chunk_size"]), (0.25, 0.2, 0.6, 7, 240))
        self.assertNotIn("speed_factor", body)

    def test_sentence_split_keeps_abbreviations(self):
        s = voice.sentences('One day Dr. Alexis Carrel sent word. "The world," he said. E. Stanley Jones came.')
        self.assertEqual(len(s), 3)

    def test_script_wraps_inserted_parts_when_switches_on(self):
        c = self.cfg()
        c["summaries"].update(chapter_preview=True, chapter_summary=True, concept_summary=True)
        ch = {"number": 1, "title": "We Go Seeking", "switches": {}}
        segs = script.chapter_script(c, {"summaries": {}}, ch, ["Para one is here.", "Para two is here.", "Para three is here."])
        kinds = [s["kind"] for s in segs]
        self.assertEqual(kinds[:5], ["header", "cue", "label", "chapter_preview", "cue"])
        self.assertEqual(kinds[-4:], ["cue", "label", "chapter_summary", "cue"])
        self.assertEqual(kinds.count("concept_summary"), 1)
        self.assertEqual([s["text"] for s in segs if s["kind"] == "author"], ["Para one is here.", "Para two is here.", "Para three is here."])

    def test_switches_default_off_and_chapter_wins(self):
        ch = {"number": 1, "title": "x", "switches": {"chapter_summary": True}}
        self.assertFalse(script.switch(CFG, {"summaries": {}}, {"switches": {}}, "chapter_summary"))
        self.assertTrue(script.switch(CFG, {"summaries": {"chapter_summary": False}}, ch, "chapter_summary"))


if __name__ == "__main__":
    unittest.main()