ASR-demo/realtime_asr_optimization_demo/tests/test_speaker_assembler.py

"""片段幂等更新和说话人安全规则的回归测试。"""

from __future__ import annotations

import unittest

from speaker_assembler import SegmentAssembler


class SegmentAssemblerTests(unittest.TestCase):
    def test_late_speaker_updates_merge_and_split_existing_blocks(self):
        """模拟 final 先到、声纹后到以及身份修正，按当前完整状态重新分块。"""
        assembler = SegmentAssembler()
        for sid in range(3):
            assembler.apply_sentence({"sentence_id": sid, "sentence": str(sid), "sentence_type": 1,
                                      "start_time": sid * 2000, "end_time": sid * 2000 + 1000})
        for sid in range(3):
            assembler.apply_speaker_update({"sentence_id": sid, "speaker_id": 0, "speaker_name": "A",
                                            "speaker_evidence": "fresh", "speaker_confidence": 0.9})
        self.assertEqual(assembler.display_blocks()[0]["segment_ids"], [0, 1, 2])
        assembler.apply_speaker_update({"sentence_id": 1, "speaker_id": 1, "speaker_name": "B",
                                        "speaker_evidence": "fresh", "speaker_confidence": 0.9})
        self.assertEqual([b["speaker_id"] for b in assembler.display_blocks()], [0, 1, 0])
        self.assertEqual(len(assembler.display_blocks(False)), 3)

    def test_speaker_update_cannot_inherit_or_change_text_and_timestamps(self):
        """异步更新同样要拒绝继承证据，且模型响应不能改写原始转写。"""
        assembler = SegmentAssembler()
        assembler.apply_sentence({"sentence_id": 0, "sentence": "原文", "start_time": 0, "end_time": 500})
        updated = assembler.apply_speaker_update({
            "sentence_id": 0, "sentence": "错误原文", "start_time": 9000,
            "speaker_id": 7, "speaker_name": "前一位", "user_id": "person-a",
            "speaker_strategy": "short_attach", "speaker_evidence": "confirmed", "speaker_confidence": 0.99,
            "_embedding": [1, 0],
        })
        self.assertEqual(updated["speaker_id"], -1)
        self.assertEqual(updated["sentence"], "原文")
        self.assertEqual(updated["start_time"], 0)
        self.assertNotIn("user_id", updated)
        self.assertNotIn("_embedding", updated)

    def test_nonfinite_or_weak_confidence_is_not_trusted(self):
        """NaN/Infinity 和不充分证据不得进入可信展示块。"""
        for confidence in (float("nan"), float("inf"), 0.59):
            assembler = SegmentAssembler()
            assembler.apply_sentence({"sentence_id": 0, "sentence": "test"})
            result = assembler.apply_speaker_update({"sentence_id": 0, "speaker_id": 0,
                "speaker_evidence": "fresh", "speaker_confidence": confidence})
            self.assertEqual(result["speaker_status"], "evidence_rejected")
            self.assertEqual(assembler.display_blocks()[0]["speaker_id"], -1)

    def test_same_cluster_different_named_identity_does_not_merge(self):
        """相同匿名簇 ID 不代表相同注册身份，弱匿名段不能挤入实名块。"""
        assembler = SegmentAssembler()
        for sid, user in enumerate(("A", "B", None)):
            assembler.apply_sentence({"sentence_id": sid, "sentence": str(sid), "speaker_id": 0,
                "speaker_evidence": "fresh", "speaker_confidence": 0.9, "user_id": user})
        self.assertEqual(len(assembler.display_blocks()), 3)

    def test_text_only_update_preserves_identity_and_final_is_not_rolled_back(self):
        """文本与声纹独立更新；重发的旧 partial 不得回滚 final。"""
        assembler = SegmentAssembler()
        assembler.apply_sentence({"sentence_id": 0, "sentence": "final", "sentence_type": 1})
        assembler.apply_speaker_update({"sentence_id": 0, "speaker_id": 0,
            "speaker_evidence": "fresh", "speaker_confidence": 0.9})
        assembler.apply_sentence({"sentence_id": 0, "sentence": "corrected", "sentence_type": 1})
        assembler.apply_sentence({"sentence_id": 0, "sentence": "old", "sentence_type": 0})
        self.assertEqual(assembler.raw_snapshot()[0]["sentence"], "corrected")
        self.assertEqual(assembler.display_blocks()[0]["speaker_id"], 0)

    def test_same_sentence_id_is_updated_in_place(self) -> None:
        assembler = SegmentAssembler()
        assembler.apply_sentence({"sentence_id": 3, "sentence": "你好", "sentence_type": 0})
        final = assembler.apply_sentence({"sentence_id": 3, "sentence": "你好，今天开始。", "sentence_type": 1})

        self.assertEqual(len(assembler.raw_snapshot()), 1)
        self.assertEqual(final["sentence_type"], 1)
        self.assertEqual(final["revision_count"], 1)

    def test_short_inherited_name_is_pending_and_embedding_is_dropped(self) -> None:
        assembler = SegmentAssembler()
        segment = assembler.apply_sentence(
            {
                "sentence_id": 1,
                "sentence": "嗯",
                "start_time": 0,
                "end_time": 900,
                "speaker_id": 7,
                "speaker_name": "上一位",
                "speaker_strategy": "short_attach",
                "speaker_evidence": "confirmed",
                "speaker_confidence": 0.99,
                "_embedding": [1, 2, 3],
            }
        )

        self.assertEqual(segment["speaker_id"], -1)
        self.assertEqual(segment["speaker_evidence"], "pending")
        self.assertNotIn("_embedding", segment)

    def test_a_b_a_keeps_three_time_ordered_blocks(self) -> None:
        assembler = SegmentAssembler()
        for sentence_id, speaker_id, start in ((1, 10, 0), (2, 20, 2000), (3, 10, 4000)):
            assembler.apply_sentence(
                {
                    "sentence_id": sentence_id,
                    "sentence": f"句子{sentence_id}",
                    "start_time": start,
                    "end_time": start + 1000,
                    "speaker_id": speaker_id,
                    "speaker_evidence": "confirmed",
                    "speaker_confidence": 0.9,
                }
            )

        blocks = assembler.display_blocks()
        self.assertEqual([block["speaker_id"] for block in blocks], [10, 20, 10])


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