"""Transcribing a file: splitting it, stamping it, and writing subtitles. to_srt is the awkward one. The text is the authority on wording and the segments on timing, and they meet at a whole-second stamp that a cleanup model was asked to leave alone. It has to survive a model that wrapped a line, dropped one, or made up a stamp nobody recorded. """ import contextlib import time import unittest import wave from unittest import mock import api import filetranscribe as ft from tests.support import DikteTest, make_wav, silence, tone class Timestamps(unittest.TestCase): def test_under_an_hour(self): self.assertEqual(ft.format_timestamp(0), "00:00") self.assertEqual(ft.format_timestamp(65.9), "01:05") self.assertEqual(ft.format_timestamp(599), "09:59") def test_past_an_hour_the_hours_show(self): self.assertEqual(ft.format_timestamp(3600), "1:00:00") self.assertEqual(ft.format_timestamp(3725), "1:02:05") def test_srt_wants_milliseconds_and_a_comma(self): self.assertEqual(ft.srt_timestamp(0), "00:00:00,000") self.assertEqual(ft.srt_timestamp(1.5), "00:00:01,500") self.assertEqual(ft.srt_timestamp(3725.25), "01:02:05,250") def test_a_negative_start_is_pulled_up_to_zero(self): self.assertEqual(ft.srt_timestamp(-3), "00:00:00,000") class ToSrt(unittest.TestCase): def test_nothing_to_do(self): self.assertEqual(ft.to_srt("", []), "") self.assertEqual(ft.to_srt("no stamps here", []), "") def test_a_cue_takes_its_timing_from_the_segment_it_came_from(self): srt = ft.to_srt("[00:01] Hello there.", [(1.25, 2.75, "hello there")]) self.assertIn("00:00:01,250 --> 00:00:02,750", srt) self.assertIn("Hello there.", srt) def test_the_text_wins_on_wording(self): """Cleanup edits survive; only the timing comes from the segments.""" srt = ft.to_srt("[00:01] Hello there.", [(1.0, 2.0, "uh hello uh there")]) self.assertIn("Hello there.", srt) self.assertNotIn("uh", srt) def test_cues_are_numbered_from_one(self): srt = ft.to_srt("[00:00] One\n[00:02] Two", [(0.0, 1.0, "One"), (2.0, 3.0, "Two")]) self.assertTrue(srt.startswith("1\n")) self.assertIn("\n2\n", srt) def test_a_wrapped_line_joins_the_cue_above_it(self): srt = ft.to_srt("[00:01] Hello there,\nand welcome.", [(1.0, 4.0, "hello there and welcome")]) self.assertIn("Hello there, and welcome.", srt) self.assertEqual(srt.count(" --> "), 1) def test_a_stamp_nobody_recorded_still_gets_timing(self): srt = ft.to_srt("[00:05] Invented.", []) self.assertIn("00:00:05,000 --> ", srt) def test_a_cue_with_no_end_runs_until_the_next_one(self): srt = ft.to_srt("[00:00] One\n[00:04] Two", []) self.assertIn("00:00:00,000 --> 00:00:04,000", srt) def test_the_last_cue_gets_a_minimum_length(self): srt = ft.to_srt("[00:10] Last words.", []) self.assertIn("00:00:10,000 --> 00:00:11,500", srt) def test_a_cue_is_cut_short_when_the_next_one_starts_first(self): """Whisper's end times overlap now and then; subtitles must not.""" srt = ft.to_srt("[00:00] One\n[00:02] Two", [(0.0, 9.0, "One"), (2.0, 3.0, "Two")]) self.assertIn("00:00:00,000 --> 00:00:02,000", srt) def test_blank_lines_and_empty_cues_are_dropped(self): srt = ft.to_srt("[00:00] One\n\n[00:02]\n[00:03] Three", []) self.assertEqual(srt.count(" --> "), 2) def test_the_hour_form_of_a_stamp_is_understood(self): srt = ft.to_srt("[1:02:05] Late.", []) self.assertIn("01:02:05,000", srt) def test_the_file_ends_with_a_newline(self): self.assertTrue(ft.to_srt("[00:00] One", []).endswith("\n")) class SplitText(unittest.TestCase): def test_short_text_stays_whole(self): self.assertEqual(ft.split_text("hello", False), ["hello"]) def test_a_long_transcript_is_broken_up(self): text = " ".join(["word"] * 8000) blocks = ft.split_text(text, False) self.assertGreater(len(blocks), 1) for block in blocks: self.assertLessEqual(len(block), ft.CLEANUP_CHUNK_CHARS) def test_nothing_is_lost_in_the_splitting(self): text = " ".join(f"word{index}" for index in range(4000)) self.assertEqual(" ".join(ft.split_text(text, False)), text) def test_a_timestamped_transcript_is_never_broken_mid_line(self): text = "\n".join(f"[00:{index:02d}] a line of some length here" for index in range(600)) blocks = ft.split_text(text, True) self.assertGreater(len(blocks), 1) for block in blocks: for line in block.splitlines(): self.assertTrue(line.startswith("[")) def test_a_single_line_longer_than_the_limit_is_kept_whole(self): text = "x" * (ft.CLEANUP_CHUNK_CHARS + 100) self.assertEqual(ft.split_text(text, True), [text]) class SplitWav(DikteTest): def wav(self, seconds, name="in.wav"): return make_wav(self.path(name), silence(seconds)) def test_a_short_file_is_handed_back_as_it_is(self): path = self.wav(2) self.assertEqual(ft.split_wav(path, self.root), [(path, 0.0)]) def test_a_long_file_is_cut_at_the_chunk_length(self): path = self.wav(5) with mock.patch.object(ft, "CHUNK_SECONDS", 2): chunks = ft.split_wav(path, self.root) self.assertEqual([offset for _, offset in chunks], [0, 2, 4]) def test_the_chunks_add_up_to_the_original(self): path = self.wav(5) with mock.patch.object(ft, "CHUNK_SECONDS", 2): chunks = ft.split_wav(path, self.root) total = 0 for chunk_path, _ in chunks: with contextlib.closing(wave.open(chunk_path, "rb")) as wav: total += wav.getnframes() self.assertEqual(wav.getframerate(), 16000) self.assertEqual(total, 5 * 16000) def test_the_chunks_do_not_write_over_each_other(self): path = self.wav(5) with mock.patch.object(ft, "CHUNK_SECONDS", 2): chunks = ft.split_wav(path, self.root) self.assertEqual(len({chunk for chunk, _ in chunks}), len(chunks)) class Transcriber(DikteTest): """The chain, with ffmpeg and both API calls faked.""" def setUp(self): super().setUp() self.source = make_wav(self.path("input.wav"), tone(1.0)) self.conf = self.config(openrouter_api_key="sk-or-test") def run_chain(self, timestamps=False, cleanup=False, transcript="raw text", segments=None, cleaned="clean text", fail=None): worker = ft.FileTranscriber(self.conf) done, failures, progress = [], [], [] worker.finished.connect(lambda *args: done.append(args)) worker.failed.connect(failures.append) worker.progress.connect(progress.append) def to_wav(path, workdir, aborter=None): return make_wav(self.path("converted.wav"), tone(1.0)) with mock.patch.object(ft, "_to_wav", side_effect=to_wav), \ mock.patch.object(ft.shutil, "which", return_value="/usr/bin/ffmpeg"), \ mock.patch.object(api, "transcribe", side_effect=fail or (lambda *a, **k: transcript)), \ mock.patch.object(api, "transcribe_segments", return_value=segments or [(0.0, 1.0, "raw text")]), \ mock.patch.object(api, "cleanup", return_value=cleaned) as cleanup_call: # The chain is run here rather than through start(): its signals are # emitted from the worker thread, and a queued connection would need # an event loop to deliver them. This is the same code, one frame down. worker._work(self.source, timestamps, cleanup) return done, failures, progress, cleanup_call def test_plain_text_out(self): done, failures, _, _ = self.run_chain() self.assertEqual(failures, []) self.assertEqual(done[0][0], "raw text") def test_cleanup_replaces_the_text(self): done, _, _, _ = self.run_chain(cleanup=True) self.assertEqual(done[0][0], "clean text") def test_cleanup_is_told_it_is_writing_subtitles(self): _, _, _, cleanup_call = self.run_chain(cleanup=True) prompt = cleanup_call.call_args.args[3] self.assertEqual(prompt, self.conf.cleanup_prompt(subtitles=True)) def test_timestamps_come_back_as_segments_and_as_stamped_lines(self): done, _, _, _ = self.run_chain( timestamps=True, segments=[(0.0, 1.0, "one"), (2.0, 3.0, "two")]) text, segments = done[0] self.assertEqual(text, "[00:00] one\n[00:02] two") self.assertEqual(len(segments), 2) def test_no_ffmpeg_installed(self): worker = ft.FileTranscriber(self.conf) failures = [] worker.failed.connect(failures.append) with mock.patch.object(ft.shutil, "which", return_value=None): worker._work(self.source, False, False) self.assertIn("ffmpeg", failures[0]) def test_an_api_failure_is_reported_rather_than_raised(self): def boom(*args, **kwargs): raise api.ApiError("OpenAI rejected the API key") _, failures, _, _ = self.run_chain(fail=boom) self.assertIn("rejected", failures[0]) def test_empty_text_is_not_sent_to_cleanup(self): _, _, _, cleanup_call = self.run_chain(cleanup=True, transcript="") cleanup_call.assert_not_called() def test_a_stopped_run_is_not_a_failure(self): def stopped(*args, **kwargs): raise api.Aborted done, failures, progress, _ = self.run_chain(fail=stopped) self.assertEqual(failures, []) self.assertEqual(done, []) self.assertEqual(progress[-1], "Stopped.") def test_the_request_is_handed_the_stop_to_watch(self): worker = ft.FileTranscriber(self.conf) with mock.patch.object(ft, "_to_wav", side_effect=lambda *a: self.source), \ mock.patch.object(ft.shutil, "which", return_value="/usr/bin/ffmpeg"), \ mock.patch.object(api, "transcribe", return_value="text") as call: worker._work(self.source, False, False) self.assertIs(call.call_args.kwargs["aborter"], worker._abort) def test_stopping_a_local_run_stops_the_model_with_it(self): """Closing the socket is nothing to a process of ours: it would grind on to the end of the chunk with nobody left to hand the answer to.""" worker = ft.FileTranscriber(self.conf) worker._local = mock.Mock() worker.stop() self.assertTrue(worker._abort.aborted) for _ in range(100): if worker._local.stop.called: break time.sleep(0.01) worker._local.stop.assert_called_once_with() def test_a_run_that_is_over_leaves_the_model_alone(self): worker = ft.FileTranscriber(self.conf) worker.stop() self.assertTrue(worker._abort.aborted) def test_a_second_start_while_one_is_running_is_ignored(self): worker = ft.FileTranscriber(self.conf) worker._thread = mock.Mock(is_alive=lambda: True) self.assertTrue(worker.busy) worker.start(self.source, False, False) self.assertTrue(worker.busy) if __name__ == "__main__": unittest.main()