whisper-stt-chatbot/test_transcription.py

import ast
import re
import unicodedata
import unittest
from pathlib import Path
from types import SimpleNamespace
from threading import Lock
from unittest.mock import Mock


class TranscriptionTests(unittest.TestCase):
    def test_artifacts_are_removed_without_losing_real_speech(self):
        tree = ast.parse(Path('app.py').read_text(encoding='utf-8'))
        functions = [node for node in tree.body if isinstance(node, ast.FunctionDef)
                     and node.name in ('is_subtitle_artifact', 'transcribe_audio')]
        model = Mock()
        texts = ["Sous-titrage ST' 501", 'Bonjour Satoko', 'Sous-titres MFP',
                 'Peux-tu afficher les sous-titres ?', 'Sous-titres réalisés par la communauté d’Amara.org']
        model.transcribe.return_value = ([SimpleNamespace(text=t) for t in texts], SimpleNamespace(language='fr'))
        ns = dict(re=re, unicodedata=unicodedata, model=model, model_lock=Lock(),
                  VAD_ENABLED=True, BEAM_SIZE=3, HOTWORDS='Satoko', logger=Mock())
        exec(compile(ast.Module(body=functions, type_ignores=[]), '<stt>', 'exec'), ns)
        self.assertEqual(ns['transcribe_audio'](None),
                         ('Bonjour Satoko Peux-tu afficher les sous-titres ?', 'fr'))
        self.assertFalse(ns['is_subtitle_artifact']('Explique le sous-titrage ST 501'))


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