old-books/_pipeline/tests/test_process.py

"""Tests for the one-command `process`. A local stub stands in for the voice server (synthetic TEST audio only).
Run: python3 -m unittest discover -s _pipeline/tests -v   (from the repo root)"""
import contextlib
import http.server
import io
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__)), ".."))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "tools"))

import audiomeasure as am
import process
import settings
import voice

try:
    import make_test_pdf
except ImportError:
    make_test_pdf = None


class Stub(http.server.BaseHTTPRequestHandler):
    posts = []
    loaded = True

    def _json(self, obj):
        data = json.dumps(obj).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 do_GET(self):
        if self.path == "/api/model-info":
            self._json({"loaded": Stub.loaded, "class_name": "StubModel", "device": "cpu"})
        elif self.path == "/get_predefined_voices":
            self._json([{"display_name": "Michael", "filename": "Michael.wav"}])
        else:
            self.send_error(404)

    def do_POST(self):
        n = int(self.headers.get("Content-Length", 0))
        body = self.rfile.read(n)
        Stub.posts.append(self.path)
        if self.path == "/restart_server":
            Stub.loaded = True
            return self._json({"message": "ok"})
        with tempfile.TemporaryDirectory() as d:
            p = os.path.join(d, "x.wav")
            expr = "0.1*(0.6*sin(2*PI*130*t)+0.3*sin(4*PI*130*t))*(0.55+0.45*sin(2*PI*4*t))"
            am.ffmpeg(["-f", "lavfi", "-i", "aevalsrc='%s':s=24000:d=10" % expr, "-ar", "24000", "-ac", "1", p])
            data = open(p, "rb").read()
        self.send_response(200)
        self.send_header("Content-Type", "audio/wav")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def log_message(self, *a):
        pass


def run_process(cfg, pdf, **kw):
    buf = io.StringIO()
    with contextlib.redirect_stdout(buf):
        rc = process.process(cfg, pdf, **kw)
    return rc, buf.getvalue()


@unittest.skipIf(make_test_pdf is None, "reportlab not installed (only needed to build the synthetic test PDF)")
class ProcessTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.server = http.server.HTTPServer(("127.0.0.1", 0), Stub)
        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 setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        d = self.tmp.name
        self.pdf = os.path.join(d, "Pipeline Smoke Test Book - by Ada Testwell.pdf")
        make_test_pdf.build(self.pdf)
        self.cfg = settings.load()
        self.cfg["books_dir"] = os.path.join(d, "books")
        self.cfg["work_dir"] = os.path.join(d, "work")
        self.cfg["voice"]["endpoint"] = self.url + "/tts"
        self.cfg["voice"]["retries"] = 0
        self.cfg["ocr"]["required"] = False          # OCR itself is proven in the CI job; here only the other stages
        Stub.posts.clear()
        Stub.loaded = True

    def tearDown(self):
        self.tmp.cleanup()

    def book_dir(self):
        return os.path.join(self.cfg["books_dir"], "testwell-ada", "pipeline-smoke-test-book")

    def test_full_run_idempotent_and_force(self):
        rc, out = run_process(self.cfg, self.pdf)
        print("\nTEST first run rc=%d" % rc)
        self.assertEqual(rc, 0, out)
        b = self.book_dir()
        names = sorted(os.listdir(os.path.join(b, "audio")))
        self.assertEqual(names, ["chapter-01.mp3", "chapter-02.mp3", "closing-credits.mp3", "opening-credits.mp3", "retail-sample.mp3", "run.log"])
        self.assertIn("OVERALL: PASS", out)
        for rel in ("text/chapter-01.txt", "text/chapter-02.txt", "script/chapter-02.json", "figures/skipped_figures.md",
                    "book.json", "warnings.log", "LICENSE_STATUS.md"):
            self.assertTrue(os.path.exists(os.path.join(b, rel)), rel)
        self.assertIn("public domain", open(os.path.join(b, "LICENSE_STATUS.md")).read())
        mtimes = {n: os.path.getmtime(os.path.join(b, "audio", n)) for n in names if n.endswith(".mp3")}
        posts = len(Stub.posts)
        self.assertGreater(posts, 0)
        rc, out = run_process(self.cfg, self.pdf)                 # second run: nothing to render
        print("TEST second run rc=%d, new voice requests: %d" % (rc, len(Stub.posts) - posts))
        self.assertEqual(rc, 0, out)
        self.assertEqual(len(Stub.posts), posts)
        self.assertIn("nothing to render", out)
        self.assertEqual(mtimes, {n: os.path.getmtime(os.path.join(b, "audio", n)) for n in mtimes})
        rc, out = run_process(self.cfg, self.pdf, force=True)     # forced: renders again
        self.assertEqual(rc, 0, out)
        self.assertGreater(len(Stub.posts), posts)

    def test_changed_script_rerenders_only_that_file(self):
        self.assertEqual(run_process(self.cfg, self.pdf)[0], 0)
        b = self.book_dir()
        path = os.path.join(b, "text", "chapter-01.txt")
        # an edit to the cleaned text is overwritten by the text stages, so change the voice setting for one run instead
        self.cfg["voice"]["extra"]["seed"] = 8
        posts = len(Stub.posts)
        rc, out = run_process(self.cfg, self.pdf)
        self.assertEqual(rc, 0, out)
        self.assertGreater(len(Stub.posts), posts)
        self.assertNotIn("skip chapter-01", out)

    def test_dry_run_renders_nothing_and_checks_voice(self):
        rc, out = run_process(self.cfg, self.pdf, dry_run=True)
        self.assertEqual(rc, 0, out)
        self.assertIn("voice server reachable", out)
        self.assertEqual(Stub.posts, [])
        self.assertFalse(os.path.exists(os.path.join(self.book_dir(), "audio")))

    def test_unreachable_voice_fails_clearly_with_no_audio(self):
        self.cfg["voice"]["endpoint"] = "http://127.0.0.1:9/tts"
        rc, out = run_process(self.cfg, self.pdf)
        print("\nTEST unreachable voice rc=%d: %s" % (rc, out.strip().splitlines()[-1]))
        self.assertEqual(rc, 1)
        self.assertIn("PROCESS FAILED: voice: voice server not reachable", out)
        self.assertFalse(os.path.exists(os.path.join(self.book_dir(), "audio")))

    def test_unloaded_model_is_loaded_on_demand(self):
        Stub.loaded = False
        rc, out = run_process(self.cfg, self.pdf)
        self.assertEqual(rc, 0, out)
        self.assertEqual(Stub.posts[0], "/restart_server")

    def test_speed_factor_is_refused(self):
        self.cfg["voice"]["extra"]["speed_factor"] = 0.9
        with self.assertRaises(voice.VoiceError):
            voice.speak(self.cfg, "hello")

    def test_ocr_required_fails_without_ocr(self):
        self.cfg["ocr"]["required"] = True
        self.cfg["ocr"]["enabled"] = False
        rc, out = run_process(self.cfg, self.pdf)
        self.assertEqual(rc, 1)
        self.assertIn("have no text layer", out)


class SmallTests(unittest.TestCase):
    def test_filename_guess(self):
        g = process.guess_from_filename
        self.assertEqual(g("/x/The Man Who Tapped - by Glenn Clark (Walter Russell).pdf"), ("The Man Who Tapped", "Glenn Clark"))
        self.assertEqual(g("/x/Some Title by Jane Doe.pdf"), ("Some Title", "Jane Doe"))
        self.assertEqual(g("/x/scan0001.pdf"), (None, None))

    def test_env_overrides(self):
        cfg = settings.load()
        settings.apply_env(cfg, {"AUDIOBOOK_VOICE_ENDPOINT": "http://1.2.3.4:8090/tts", "AUDIOBOOK_REQUIRE_OCR": "0"})
        self.assertEqual(cfg["voice"]["endpoint"], "http://1.2.3.4:8090/tts")
        self.assertFalse(cfg["ocr"]["required"])

    def test_locked_voice_settings_unchanged(self):
        v = settings.load()["voice"]
        self.assertEqual((v["voice"], v["extra"]["exaggeration"], v["extra"]["cfg_weight"], v["extra"]["temperature"],
                          v["extra"]["seed"], v["extra"]["split_text"], v["extra"]["chunk_size"]), ("Michael.wav", 0.25, 0.2, 0.6, 7, True, 240))
        self.assertNotIn("speed_factor", v["extra"])


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