248 lines
8.6 KiB
Python
248 lines
8.6 KiB
Python
# Copyright 2026 Umbra Tools (umbratools.com)
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
|
|
from umbraaudio.analysis.name_tokens import (
|
|
_is_known,
|
|
_is_role_token,
|
|
_normalize_name_token,
|
|
_soundex,
|
|
ensure_name_tokens_table,
|
|
merge_tokens,
|
|
persist_name_tokens,
|
|
scan_utterances,
|
|
suggest_speaker,
|
|
)
|
|
|
|
ALIAS_MAP = {
|
|
"tobias": "Tobias",
|
|
"teppo": "Teppo",
|
|
"kainda": "Kainda",
|
|
"kainda walters foster murphy": "Kainda Walters Foster Murphy",
|
|
"chief of the guard": "Chief of the Guard",
|
|
"the king": "The King",
|
|
"narrator": "Narrator",
|
|
}
|
|
|
|
|
|
def _u(idx: int, text: str, typ: str, speaker: str = "narrator") -> dict:
|
|
return {
|
|
"utterance_index": idx,
|
|
"text": text,
|
|
"type": typ,
|
|
"speaker_name": speaker,
|
|
}
|
|
|
|
|
|
class TestNormalization:
|
|
def test_scene_break_contamination(self) -> None:
|
|
assert _normalize_name_token("At The Gate\nTobias") == "Tobias"
|
|
|
|
def test_gerund_leadin(self) -> None:
|
|
assert _normalize_name_token("Nodding Tobias") == "Tobias"
|
|
|
|
def test_temporal_leadin(self) -> None:
|
|
assert _normalize_name_token("Now Milfred") == "Milfred"
|
|
|
|
def test_role_phrase_preserved(self) -> None:
|
|
# The determiner is part of the role reference; keep it intact.
|
|
assert _normalize_name_token("The first guard") == "The first guard"
|
|
assert _normalize_name_token("Royal Guard") == "Royal Guard"
|
|
|
|
def test_bare_determiner_dropped(self) -> None:
|
|
assert _normalize_name_token("The") == ""
|
|
|
|
def test_indefinite_pronoun_dropped(self) -> None:
|
|
assert _normalize_name_token("Another") == ""
|
|
assert _normalize_name_token("One") == ""
|
|
|
|
def test_conjoined_names_kept(self) -> None:
|
|
# Two adjacent speakers in a tag: keep the full phrase for review.
|
|
assert _normalize_name_token("Tobias Teppo") == "Tobias Teppo"
|
|
|
|
def test_empty_input(self) -> None:
|
|
assert _normalize_name_token("") == ""
|
|
assert _normalize_name_token("...") == ""
|
|
|
|
|
|
class TestRoleDetection:
|
|
def test_role_nouns(self) -> None:
|
|
for t in ["The guard", "Royal Guard", "The Kid", "Advisor", "The officer"]:
|
|
assert _is_role_token(t) is True
|
|
|
|
def test_personal_names_not_roles(self) -> None:
|
|
for t in ["Tobias", "Kainda", "Tobias Teppo", "Nering"]:
|
|
assert _is_role_token(t) is False
|
|
|
|
|
|
class TestIsKnown:
|
|
def test_exact(self) -> None:
|
|
assert _is_known("Tobias", ALIAS_MAP) is True
|
|
|
|
def test_containment_distinctive(self) -> None:
|
|
# "Kainda" is a distinctive name contained in a longer profile.
|
|
assert _is_known("Kainda", ALIAS_MAP) is True
|
|
|
|
def test_single_role_word_not_known(self) -> None:
|
|
# "guard" alone is a DIFFERENT speaker from "Chief of the Guard".
|
|
assert _is_known("guard", ALIAS_MAP) is False
|
|
assert _is_known("voice", ALIAS_MAP) is False
|
|
|
|
def test_single_identity_word_not_absorbed_by_longer_profile(self) -> None:
|
|
# Regression guard: an unnamed "Witch" must NOT be auto-resolved into a
|
|
# distinct profile named "The Witch Pair" via word containment.
|
|
amap = {"the witch pair": "The Witch Pair"}
|
|
assert _is_known("Witch", amap) is False
|
|
# ...but the distinctive multi-word name still resolves.
|
|
assert _is_known("The Witch Pair", amap) is True
|
|
|
|
def test_empty_always_known(self) -> None:
|
|
assert _is_known("", ALIAS_MAP) is True
|
|
|
|
|
|
class TestSuggestSpeaker:
|
|
def test_soundex_misspelling(self) -> None:
|
|
name, score = suggest_speaker("Nering", {"neriah": "Neriah"})
|
|
assert name == "Neriah"
|
|
assert score >= 0.72
|
|
|
|
def test_close_spelling_variant(self) -> None:
|
|
# "Nering"/"Neriah" share a first letter and one small edit.
|
|
name, _score = suggest_speaker("Nering", {"neriah": "Neriah"})
|
|
assert name == "Neriah"
|
|
|
|
def test_containment_suggestion(self) -> None:
|
|
name, score = suggest_speaker("Tobias", {"tobias teppo": "Tobias Teppo"})
|
|
assert name == "Tobias Teppo"
|
|
assert score >= 0.72
|
|
|
|
def test_no_suggestion_for_unrelated(self) -> None:
|
|
name, _score = suggest_speaker("Zorblax", {"tobias": "Tobias"})
|
|
assert name is None
|
|
|
|
def test_article_not_soundex_matched(self) -> None:
|
|
# "The first guard" should NOT match "The General"/"The King" merely on
|
|
# the shared article via soundex.
|
|
amap = {"the first guard": "The First Guard", "the king": "The King"}
|
|
name, _ = suggest_speaker("The first guard", amap)
|
|
assert name == "The First Guard"
|
|
|
|
|
|
class TestSoundex:
|
|
def test_known_codes(self) -> None:
|
|
assert _soundex("Robert") == "R163"
|
|
assert _soundex("Rupert") == "R163"
|
|
assert _soundex("Tymczak") == "T522"
|
|
assert _soundex("Pfister") == "P236"
|
|
|
|
def test_empty(self) -> None:
|
|
assert _soundex("") == ""
|
|
|
|
|
|
class TestScanUtterances:
|
|
def test_name_tag_orphan(self) -> None:
|
|
utts = [
|
|
_u(0, "Nering said,", "narration"),
|
|
_u(1, "We must go now.", "dialogue", "Nering"),
|
|
]
|
|
toks = scan_utterances(utts, 3, ALIAS_MAP)
|
|
tokens = {t.token for t in toks}
|
|
assert "Nering" in tokens
|
|
|
|
def test_speaker_name_orphan_raw(self) -> None:
|
|
utts = [_u(0, "Stop there.", "dialogue", "The first guard")]
|
|
toks = scan_utterances(utts, 3, ALIAS_MAP)
|
|
# The raw stored speaker_name is preserved (assign matches it exactly).
|
|
assert [t.token for t in toks] == ["The first guard"]
|
|
assert toks[0].kind == "speaker_name"
|
|
|
|
def test_known_speaker_not_flagged(self) -> None:
|
|
utts = [
|
|
_u(0, "Kainda says,", "narration"),
|
|
_u(1, "We are ready.", "dialogue", "Kainda"),
|
|
]
|
|
assert scan_utterances(utts, 3, ALIAS_MAP) == []
|
|
|
|
def test_narrator_dialogue_not_flagged(self) -> None:
|
|
utts = [_u(0, "Hello there.", "dialogue", "narrator")]
|
|
assert scan_utterances(utts, 3, ALIAS_MAP) == []
|
|
|
|
def test_role_token_has_no_suggestion(self) -> None:
|
|
utts = [_u(0, "Move.", "dialogue", "The Kid")]
|
|
toks = scan_utterances(utts, 3, ALIAS_MAP)
|
|
assert toks[0].suggestion is None
|
|
|
|
|
|
class TestMergeTokens:
|
|
def test_merges_across_chapters(self) -> None:
|
|
c1 = [_u(0, "Stop.", "dialogue", "Advisor")]
|
|
c2 = [_u(0, "Wait.", "dialogue", "Advisor")]
|
|
merged = merge_tokens(
|
|
[
|
|
scan_utterances(c1, 1, ALIAS_MAP),
|
|
scan_utterances(c2, 2, ALIAS_MAP),
|
|
],
|
|
ALIAS_MAP,
|
|
)
|
|
assert len(merged) == 1
|
|
assert merged[0].occurrences == 2
|
|
assert merged[0].chapters == [1, 2]
|
|
|
|
def test_sorted_by_occurrences(self) -> None:
|
|
per = [
|
|
scan_utterances(
|
|
[
|
|
_u(0, "A.", "dialogue", "Advisor"),
|
|
_u(1, "B.", "dialogue", "Royal Guard"),
|
|
_u(2, "C.", "dialogue", "Royal Guard"),
|
|
],
|
|
1,
|
|
ALIAS_MAP,
|
|
)
|
|
]
|
|
merged = merge_tokens(per, ALIAS_MAP)
|
|
assert merged[0].token == "Royal Guard"
|
|
|
|
|
|
class TestPersistence:
|
|
def _db(self) -> sqlite3.Connection:
|
|
db = sqlite3.connect(":memory:")
|
|
db.row_factory = sqlite3.Row
|
|
db.execute("CREATE TABLE project (id INTEGER PRIMARY KEY)")
|
|
db.execute("INSERT INTO project (id) VALUES (1)")
|
|
ensure_name_tokens_table(db)
|
|
return db
|
|
|
|
def test_insert_and_preserve_status(self) -> None:
|
|
db = self._db()
|
|
utts = [_u(0, "Stop.", "dialogue", "Advisor")]
|
|
tokens = scan_utterances(utts, 1, ALIAS_MAP)
|
|
inserted, updated = persist_name_tokens(db, 1, tokens)
|
|
assert inserted == 1 and updated == 0
|
|
|
|
# Mark it dismissed, then re-persist the same token.
|
|
db.execute(
|
|
"UPDATE name_tokens SET status='dismissed', target_speaker=NULL "
|
|
"WHERE token='Advisor'"
|
|
)
|
|
db.commit()
|
|
inserted2, updated2 = persist_name_tokens(db, 1, tokens)
|
|
assert inserted2 == 0 and updated2 == 1
|
|
row = db.execute("SELECT status FROM name_tokens WHERE token='Advisor'").fetchone()
|
|
assert row["status"] == "dismissed" # human decision preserved
|
|
db.close()
|
|
|
|
def test_idempotent(self) -> None:
|
|
db = self._db()
|
|
tokens = scan_utterances(
|
|
[_u(0, "Stop.", "dialogue", "Advisor")], 1, ALIAS_MAP
|
|
)
|
|
persist_name_tokens(db, 1, tokens)
|
|
persist_name_tokens(db, 1, tokens)
|
|
count = db.execute("SELECT COUNT(*) c FROM name_tokens").fetchone()["c"]
|
|
assert count == 1
|
|
db.close()
|