tts-voice-creator-clone-and-design-2/tests/test_attribution_spans.py

"""Span-based speaker attribution: request validation, prompt, parsing, route."""
from __future__ import annotations

import os
import tempfile
import unittest
from unittest import mock

os.environ.setdefault("CONFIG_DIR", tempfile.mkdtemp(prefix="ttvc-spans-test-config-"))

from fastapi import FastAPI
from fastapi.testclient import TestClient

from core.attribution_spans import (
    SpanRequestError, build_messages, max_tokens_for, parse_response, validate_spans,
)
import routes.conversation as conversation

SPANS = [
    {"id": "N0", "kind": "narration", "text": "Der Zwerg schneuzte sich."},
    {"id": "D1", "kind": "dialogue", "text": "Na los, steh hier nicht weiter rum."},
    {"id": "N2", "kind": "narration", "text": "sagte Linosch."},
    {"id": "D3", "kind": "dialogue", "text": "Vielen Dank!"},
]


class ValidateSpansTests(unittest.TestCase):
    def test_accepts_a_well_formed_window(self):
        self.assertEqual([s["id"] for s in validate_spans(SPANS)], ["N0", "D1", "N2", "D3"])

    def test_rejects_bad_shapes(self):
        bad = [
            None, [], "x", [1],
            [{"id": "X1", "kind": "dialogue", "text": "a"}],                    # id pattern
            [{"id": "D1", "kind": "narration", "text": "a"}],                   # kind/id mismatch
            [{"id": "D1", "kind": "dialogue", "text": "  "}],                   # empty text
            [{"id": "D1", "kind": "dialogue", "text": "a"}, {"id": "D1", "kind": "dialogue", "text": "b"}],
            [{"id": "N1", "kind": "narration", "text": "only narration"}],      # nothing to attribute
            [{"id": "D%d" % i, "kind": "dialogue", "text": "a"} for i in range(401)],
            [{"id": "D1", "kind": "dialogue", "text": "x" * 20_001}],
        ]
        for case in bad:
            with self.subTest(case=str(case)[:60]):
                with self.assertRaises(SpanRequestError):
                    validate_spans(case)


class PromptTests(unittest.TestCase):
    def test_prompt_lists_every_line_and_the_ids_to_attribute(self):
        system, user = build_messages(SPANS, ["Marcian"], "Marcian: Hallo", "Danach ging er.", "German")
        self.assertIn("Do NOT rewrite", system)
        self.assertIn("German", system)
        self.assertIn('[D1] "Na los, steh hier nicht weiter rum."', user)
        self.assertIn("[N2] sagte Linosch.", user)
        self.assertIn("Known characters so far: Marcian", user)
        self.assertIn("Recent dialogue", user)
        self.assertIn("Attribute these lines: D1, D3", user)

    def test_only_ids_restricts_the_request_for_a_retry(self):
        _, user = build_messages(SPANS, only_ids=["D3"])
        self.assertIn("Attribute these lines: D3", user)
        self.assertNotIn("D1, D3", user)

    def test_token_budget_grows_with_dialogue_and_is_capped(self):
        many = lambda n: [{"id": "D%d" % i, "kind": "dialogue", "text": "a"} for i in range(n)]
        self.assertLess(max_tokens_for(SPANS), max_tokens_for(many(30)))
        self.assertLessEqual(max_tokens_for(many(400)), 4096)


class ParseTests(unittest.TestCase):
    def test_parses_clean_json_and_normalises_names(self):
        raw = '{"speakers":[{"id":"D1","speaker":" Linosch ","emotion":"gereizt"},{"id":"D3","speaker":"unbekannt","emotion":""}]}'
        out = parse_response(raw, validate_spans(SPANS))
        self.assertEqual(out["speakers"]["D1"], {"speaker": "Linosch", "emotion": "gereizt"})
        self.assertEqual(out["speakers"]["D3"]["speaker"], "Unknown")
        self.assertEqual(out["missing"], [])
        self.assertEqual(out["characters"], ["Linosch"])

    def test_ignores_reasoning_fences_and_unknown_ids(self):
        raw = ('<think>D1 is Linosch, {"id":"D1","speaker":"Wrong"}</think>\n```json\n'
               '{"speakers":[{"id":"D1","speaker":"Linosch","emotion":""},{"id":"D9","speaker":"Ghost"}]}\n```')
        out = parse_response(raw, validate_spans(SPANS))
        self.assertEqual(out["speakers"]["D1"]["speaker"], "Linosch")
        self.assertNotIn("D9", out["speakers"])
        self.assertEqual(out["missing"], ["D3"])

    def test_salvages_a_truncated_answer(self):
        raw = '{"speakers":[{"id":"D1","speaker":"Linosch","emotion":"x"},{"id":"D3","speak'
        out = parse_response(raw, validate_spans(SPANS))
        self.assertEqual(list(out["speakers"]), ["D1"])
        self.assertEqual(out["missing"], ["D3"])

    def test_narrator_answer_for_a_non_speech_line_is_kept_out_of_the_cast(self):
        raw = '{"speakers":[{"id":"D1","speaker":"Erzähler"},{"id":"D3","speaker":"Marcian"}]}'
        out = parse_response(raw, validate_spans(SPANS))
        self.assertEqual(out["speakers"]["D1"]["speaker"], "Narrator")
        self.assertEqual(out["characters"], ["Marcian"])


class RouteTests(unittest.TestCase):
    def setUp(self):
        app = FastAPI()
        app.include_router(conversation.router)
        self.client = TestClient(app)

    def test_rejects_an_invalid_request_before_calling_the_llm(self):
        with mock.patch.object(conversation, "_span_llm_answer") as llm:
            r = self.client.post("/api/attribute-spans", json={"spans": [{"id": "N1", "kind": "narration", "text": "x"}]})
        self.assertEqual(r.status_code, 400)
        llm.assert_not_called()

    def test_returns_speakers_and_never_the_text(self):
        answer = '{"speakers":[{"id":"D1","speaker":"Linosch","emotion":"gereizt"},{"id":"D3","speaker":"Alrik","emotion":""}]}'
        with mock.patch.object(conversation, "_span_llm_answer", return_value=answer) as llm:
            r = self.client.post("/api/attribute-spans", json={"spans": SPANS, "known_characters": ["Alrik"]})
        self.assertEqual(r.status_code, 200)
        body = r.json()
        self.assertEqual(body["speakers"]["D3"]["speaker"], "Alrik")
        self.assertEqual(body["missing"], [])
        self.assertNotIn("Zwerg", r.text)
        llm.assert_called_once()

    def test_retries_only_the_lines_the_model_skipped(self):
        first = '{"speakers":[{"id":"D1","speaker":"Linosch","emotion":""}]}'
        second = '{"speakers":[{"id":"D3","speaker":"Alrik","emotion":""}]}'
        with mock.patch.object(conversation, "_span_llm_answer", side_effect=[first, second]) as llm:
            r = self.client.post("/api/attribute-spans", json={"spans": SPANS})
        body = r.json()
        self.assertEqual(llm.call_count, 2)
        self.assertEqual({k: v["speaker"] for k, v in body["speakers"].items()}, {"D1": "Linosch", "D3": "Alrik"})
        self.assertEqual(body["missing"], [])
        retry_prompt = llm.call_args_list[1].args[0]["messages"][1]["content"]
        self.assertIn("Attribute these lines: D3", retry_prompt)

    def test_reports_failure_instead_of_erroring_when_the_llm_is_down(self):
        with mock.patch.object(conversation, "_span_llm_answer", side_effect=RuntimeError("LLM HTTP 503")):
            r = self.client.post("/api/attribute-spans", json={"spans": SPANS})
        self.assertEqual(r.status_code, 200)
        body = r.json()
        self.assertTrue(body["failed"])
        self.assertEqual(body["missing"], ["D1", "D3"])


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