from __future__ import annotations import json import unittest from moss_transcribe_diarize.subtitle import SubtitleSegment, SubtitleStyle, export_ass, export_json, export_srt from moss_transcribe_diarize.subtitle.export import format_ass_time, format_srt_time class SubtitleExportTest(unittest.TestCase): def test_time_formatters(self): self.assertEqual(format_srt_time(3661.234), "01:01:01,234") self.assertEqual(format_ass_time(3661.23), "1:01:01.23") def test_export_srt(self): text = export_srt([SubtitleSegment("seg_0001", 0.5, 2.0, "S01", "hello")]) self.assertIn("00:00:00,500 --> 00:00:02,000", text) self.assertIn("S01: hello", text) def test_export_srt_with_speaker_names(self): text = export_srt( [SubtitleSegment("seg_0001", 0.5, 2.0, "S01", "hello")], speaker_names={"S01": "Alice"}, ) self.assertIn("Alice: hello", text) def test_export_ass(self): text = export_ass( [SubtitleSegment("seg_0001", 0.5, 2.0, "S01", "hello")], style=SubtitleStyle(font_size=42, show_speaker=False), video_width=1280, video_height=720, ) self.assertIn("PlayResX: 1280", text) self.assertIn("Style: Speaker_S01,Noto Sans CJK SC,42", text) self.assertIn("Dialogue: 0,0:00:00.50,0:00:02.00,Speaker_S01", text) self.assertIn("hello", text) def test_export_ass_with_speaker_names(self): text = export_ass( [SubtitleSegment("seg_0001", 0.5, 2.0, "S01", "hello")], style=SubtitleStyle(font_size=42, speaker_names={"S01": "Alice"}), video_width=1280, video_height=720, ) self.assertIn("Alice: hello", text) def test_export_json(self): data = json.loads(export_json([SubtitleSegment("seg_0001", 0, 1, "S01", "hello")])) self.assertEqual(data[0]["id"], "seg_0001") self.assertEqual(data[0]["text"], "hello") if __name__ == "__main__": unittest.main()