211 lines
7.8 KiB
Python
211 lines
7.8 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright 2026 Umbra Tools (umbratools.com)
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Regression tests for resegment_merge.py."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from typing import TYPE_CHECKING
|
|
|
|
from scripts.resegment_merge import _merge_chapter, _norm, _TargetPool
|
|
|
|
if TYPE_CHECKING:
|
|
from pathlib import Path
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# helpers
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def _make_db(tmp_path: Path, chapter_rows: list[dict], title: str, source: str) -> Path:
|
|
"""Create a minimal project.db with a single chapter + utterance rows."""
|
|
db_path = tmp_path / "project.db"
|
|
db = sqlite3.connect(str(db_path))
|
|
db.execute("""
|
|
CREATE TABLE chapters (
|
|
id INTEGER PRIMARY KEY,
|
|
title TEXT,
|
|
chapter_index INTEGER NOT NULL,
|
|
source_text TEXT,
|
|
status TEXT
|
|
)
|
|
""")
|
|
db.execute("""
|
|
CREATE TABLE utterances (
|
|
id INTEGER PRIMARY KEY,
|
|
chapter_id INTEGER NOT NULL,
|
|
utterance_index INTEGER NOT NULL,
|
|
text TEXT NOT NULL,
|
|
type TEXT NOT NULL DEFAULT 'narration',
|
|
context_before TEXT,
|
|
context_after TEXT,
|
|
emotion TEXT,
|
|
pace TEXT,
|
|
emphasis_words TEXT,
|
|
pause_before_ms INTEGER,
|
|
pause_after_ms INTEGER,
|
|
instruct TEXT,
|
|
speaker_name TEXT,
|
|
attribution_confidence REAL DEFAULT 1.0,
|
|
attribution_reasoning TEXT,
|
|
emotion_confidence REAL DEFAULT 1.0,
|
|
user_edited INTEGER NOT NULL DEFAULT 0,
|
|
status TEXT NOT NULL DEFAULT 'pending'
|
|
)
|
|
""")
|
|
db.execute("INSERT INTO chapters (id, title, chapter_index, source_text) VALUES (1,?,?,?)",
|
|
(title, 1, source))
|
|
for r in chapter_rows:
|
|
db.execute(
|
|
"INSERT INTO utterances (id, chapter_id, utterance_index, text, type, "
|
|
"speaker_name, user_edited, emotion, pace, instruct, status) "
|
|
"VALUES (?, 1, ?, ?, ?, ?, ?, ?, ?, ?, 'pending')",
|
|
(r["id"], r["index"], r["text"], r["type"],
|
|
r.get("speaker", "narrator"), r.get("user_edited", 0),
|
|
r.get("emotion"), r.get("pace"), r.get("instruct")),
|
|
)
|
|
db.commit()
|
|
db.close()
|
|
return db_path
|
|
|
|
|
|
def _read_rows(db_path: Path, chapter_id: int = 1) -> list[dict]:
|
|
db = sqlite3.connect(str(db_path))
|
|
db.row_factory = sqlite3.Row
|
|
rows = db.execute(
|
|
"SELECT id, utterance_index, text, type, speaker_name, user_edited "
|
|
"FROM utterances WHERE chapter_id=? ORDER BY utterance_index",
|
|
(chapter_id,),
|
|
).fetchall()
|
|
db.close()
|
|
return [dict(r) for r in rows]
|
|
|
|
|
|
def _make_seg():
|
|
from umbraaudio.mapping.segmenter import Segmenter
|
|
return Segmenter()
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# tests
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
class TestNorm:
|
|
def test_strips_quotes_and_commas(self):
|
|
assert _norm("\u201cHello,\u201d he said.") == "hello, he said"
|
|
|
|
def test_case_insensitive(self):
|
|
assert _norm("Hallr says,") == _norm("hallr says,")
|
|
|
|
|
|
class TestTargetPool:
|
|
def test_exact_pop(self):
|
|
pool = _TargetPool([
|
|
{"text": "Hallr shrugs.", "type": "narration"},
|
|
{"text": "Hallr shrugs.", "type": "narration"},
|
|
])
|
|
s = pool.exact("hallr shrugs")
|
|
assert s is not None
|
|
assert s["text"] == "Hallr shrugs."
|
|
s2 = pool.exact("hallr shrugs")
|
|
assert s2 is not None
|
|
assert pool.exact("hallr shrugs") is None
|
|
|
|
def test_narration_rewrite_containment(self):
|
|
pool = _TargetPool([
|
|
{"text": "Ready to Fight\nHallr was up early.", "type": "narration"},
|
|
])
|
|
s = pool.narration_rewrite("hallr was up early")
|
|
assert s is not None
|
|
assert "Hallr" in s["text"]
|
|
|
|
|
|
class TestMergeChapter:
|
|
def test_dialogue_preserved(self, tmp_path):
|
|
source = 'Hallr says, \u201cI will fight.\u201d\nFeliu nods.\n\u201cSounds good.\u201d'
|
|
db = _make_db(tmp_path, [
|
|
{"id": 1, "index": 0, "text": "\u201cI will fight.\u201d",
|
|
"type": "dialogue", "speaker": "Hallr"},
|
|
{"id": 2, "index": 1, "text": "Feliu nods.", "type": "narration"},
|
|
{"id": 3, "index": 2, "text": "\u201cSounds good.\u201d",
|
|
"type": "dialogue", "speaker": "Feliu"},
|
|
], "Test", source)
|
|
seg = _make_seg()
|
|
kept, new, rewritten, dropped = _merge_chapter(
|
|
seg, db, _read_db_chapter(db, 1), ["Hallr", "Feliu"], 0.82, dry_run=True
|
|
)
|
|
rows = _read_rows(db, 1)
|
|
# dialogue ids and speakers unchanged
|
|
assert rows[0]["id"] == 1 and rows[0]["speaker_name"] == "Hallr"
|
|
assert rows[2]["id"] == 3 and rows[2]["speaker_name"] == "Feliu"
|
|
|
|
def test_leading_comma_rewrite(self, tmp_path):
|
|
source = "Hallr shrugs.\nHe moves on."
|
|
db = _make_db(tmp_path, [
|
|
{"id": 1, "index": 0, "text": ", Hallr shrugs.", "type": "narration"},
|
|
{"id": 2, "index": 1, "text": "He moves on.", "type": "narration"},
|
|
], "Test", source)
|
|
seg = _make_seg()
|
|
_merge_chapter(
|
|
seg, db, _read_db_chapter(db, 1), [], 0.82, dry_run=False
|
|
)
|
|
rows = _read_rows(db, 1)
|
|
assert rows[0]["text"] == "Hallr shrugs."
|
|
assert rows[0]["speaker_name"] == "narrator"
|
|
|
|
def test_user_edited_row_not_dropped(self, tmp_path):
|
|
source = "Hallr was tired."
|
|
db = _make_db(tmp_path, [
|
|
{"id": 1, "index": 0, "text": "Hallr was tired.", "type": "narration"},
|
|
{"id": 2, "index": 1, "text": "I hope you like the story.",
|
|
"type": "narration", "user_edited": 1},
|
|
], "Test", source)
|
|
seg = _make_seg()
|
|
kept, new, rewritten, dropped = _merge_chapter(
|
|
seg, db, _read_db_chapter(db, 1), [], 0.82, dry_run=True
|
|
)
|
|
assert dropped == []
|
|
rows = _read_rows(db, 1)
|
|
assert len(rows) == 2
|
|
|
|
def test_scene_break_preserved_and_added(self, tmp_path):
|
|
source = "Scene one.\n<<<>>>\nScene two."
|
|
db = _make_db(tmp_path, [
|
|
{"id": 1, "index": 0, "text": "Scene one.", "type": "narration"},
|
|
{"id": 2, "index": 1, "text": "<<<>>>", "type": "scene_break"},
|
|
{"id": 3, "index": 2, "text": "Scene two.", "type": "narration"},
|
|
], "Test", source)
|
|
seg = _make_seg()
|
|
_merge_chapter(
|
|
seg, db, _read_db_chapter(db, 1), [], 0.82, dry_run=False
|
|
)
|
|
rows = _read_rows(db, 1)
|
|
scene_rows = [r for r in rows if r["type"] == "scene_break"]
|
|
assert len(scene_rows) == 1
|
|
assert scene_rows[0]["speaker_name"] == "scene_break"
|
|
|
|
def test_indices_renumbered(self, tmp_path):
|
|
source = "One.\n<<<>>>\nTwo.\nThree."
|
|
db = _make_db(tmp_path, [
|
|
{"id": 1, "index": 0, "text": "One.", "type": "narration"},
|
|
{"id": 3, "index": 2, "text": "Three.", "type": "narration"},
|
|
{"id": 2, "index": 1, "text": "Two.", "type": "narration"},
|
|
], "Test", source)
|
|
seg = _make_seg()
|
|
_merge_chapter(
|
|
seg, db, _read_db_chapter(db, 1), [], 0.82, dry_run=False
|
|
)
|
|
rows = _read_rows(db, 1)
|
|
indices = [r["utterance_index"] for r in rows]
|
|
assert indices == list(range(len(rows)))
|
|
|
|
|
|
def _read_db_chapter(db_path: Path, chapter_id: int) -> dict:
|
|
db = sqlite3.connect(str(db_path))
|
|
db.row_factory = sqlite3.Row
|
|
r = db.execute("SELECT id, title, chapter_index, source_text FROM chapters WHERE id=?",
|
|
(chapter_id,)).fetchone()
|
|
db.close()
|
|
return dict(r)
|