test_shape_render.py

tests/test_shape_render.py

1. import json
3. from transcription_pipeline import render
4. from transcription_pipeline.shape import (_majority_speaker, _speaker_runs, rename_speakers,
5. sentences_from_segments, turns)
7. from .helpers import fake_result
10. def shaped(speakers=True):
11. return rename_speakers(sentences_from_segments(fake_result(speakers)["segments"]))
14. def test_sentences_take_word_times_and_speakers():
15. sents, names = shaped()
16. assert [s["text"] for s in sents] == ["Hi, thanks for calling back.", "No problem.",
17. "Did the draft come through?"]
18. assert [s["start"] for s in sents] == [0.5, 2.2, 3.5]
19. assert all(s["timed"] for s in sents)
20. # SPEAKER_01 talks first, so it becomes Speaker 1 whatever its raw label.
21. assert [s["speaker"] for s in sents] == ["Speaker 1", "Speaker 2", "Speaker 2"]
22. assert names == ["Speaker 1", "Speaker 2"]
25. def test_turns_merge_consecutive_sentences():
26. sents, _ = shaped()
27. t = turns(sents)
28. assert [x["speaker"] for x in t] == ["Speaker 1", "Speaker 2"]
29. assert len(t[1]["sentences"]) == 2 and t[1]["end"] == 5.0
32. def test_unaligned_segment_falls_back_to_interpolation():
33. seg = {"start": 10.0, "end": 20.0, "text": "First part here. Second part here."}
34. sents = sentences_from_segments([seg])
35. assert sents[0]["start"] == 10.0 and not sents[0]["timed"]
36. # "Second" starts at character 17 of 34, so halfway through the segment.
37. assert sents[1]["start"] == 15.0
40. def test_unpunctuated_speech_splits_at_speaker_change():
41. words = [{"word": x, "start": i * 0.4, "end": i * 0.4 + 0.3, "speaker": s}
42. for i, (x, s) in enumerate([("so", "A"), ("are", "A"), ("we", "A"), ("done", "A"),
43. ("yes", "B"), ("all", "B"), ("done", "B")])]
44. seg = {"start": 0.0, "end": 3.0, "text": "so are we done yes all done", "words": words}
45. sents, _ = rename_speakers(sentences_from_segments([seg]))
46. assert [(s["speaker"], s["text"]) for s in sents] == [("Speaker 1", "so are we done"),
47. ("Speaker 2", "yes all done")]
48. assert sents[1]["start"] == 1.6
51. def test_one_word_flicker_is_not_a_speaker_change():
52. labels = ["A", "A", "A", "B", "A", "A"]
53. items = [(i, {"word": "w", "speaker": s}) for i, s in enumerate(labels)]
54. runs = _speaker_runs(items)
55. assert len(runs) == 1 and runs[0][0] == "A" and len(runs[0][1]) == 6
58. def test_words_without_speaker_join_their_run():
59. items = [(0, {"word": "a"}), (1, {"word": "b", "speaker": "A"}),
60. (2, {"word": "c", "speaker": "A"}), (3, {"word": "d"}),
61. (4, {"word": "e", "speaker": "B"}), (5, {"word": "f", "speaker": "B"})]
62. runs = _speaker_runs(items)
63. assert [(r[0], len(r[1])) for r in runs] == [("A", 4), ("B", 2)]
66. def test_majority_tie_goes_to_first_speaker():
67. words = [{"word": "one", "start": 0, "end": 1, "speaker": "X"},
68. {"word": "two", "start": 1, "end": 2, "speaker": "Y"}]
69. assert _majority_speaker(words) == "X"
72. def test_no_diarization_means_no_labels():
73. sents, names = shaped(speakers=False)
74. assert names == [] and all(s["speaker"] is None for s in sents)
75. assert render.to_txt(sents).splitlines()[0] == "[00:00:00] Hi, thanks for calling back."
78. def test_txt_line_format():
79. sents, _ = shaped()
80. lines = render.to_txt(sents).splitlines()
81. assert lines[0] == "[00:00:00] Speaker 1: Hi, thanks for calling back."
82. assert lines[1] == "[00:00:02] Speaker 2: No problem."
83. bare = render.to_txt(sents, timestamps=False).splitlines()
84. assert bare[0] == "Speaker 1: Hi, thanks for calling back."
85. assert render.stamp(3725.9) == "[01:02:05]"
88. def test_markdown_is_a_vault_note():
89. sents, names = shaped()
90. md = render.to_markdown(sents, {"title": "Call", "recorded": "2026-09-01",
91. "transcribed": "2026-09-01T10:00:00", "duration": "0:00:05",
92. "engine": "local", "language": "en", "source": "call.m4a",
93. "speakers": names})
94. head, body = md.split("---\n", 2)[1:]
95. assert "type: transcript" in head and ' - "Speaker 2"' in head
96. assert "**Speaker 1** [00:00:00]\nHi, thanks for calling back." in body
97. assert "**Speaker 2** [00:00:02]\nNo problem. Did the draft come through?" in body
100. def test_subtitles():
101. sents, _ = shaped()
102. srt = render.to_srt(sents)
103. assert srt.startswith("1\n00:00:00,500 --> 00:00:01,700\n"
104. "Speaker 1: Hi, thanks for calling back.\n")
105. vtt = render.to_vtt(sents)
106. assert vtt.startswith("WEBVTT\n\n00:00:00.500 --> 00:00:01.700\n")
107. # An untimed sentence with no end still gets a cue that moves forward.
108. cue = list(render._cues([{"text": "x", "start": 4.0, "end": 4.0, "speaker": None}]))[0]
109. assert cue[1] > cue[0]
112. def test_json_carries_everything():
113. sents, names = shaped()
114. data = json.loads(render.to_json(sents, fake_result(), {"speakers": names}))
115. assert data["meta"]["speakers"] == names and len(data["segments"]) == 2
116. assert data["sentences"][0]["timed"] is True