diff --git a/tests/test_server.py b/tests/test_server.py new file mode 100644 index 0000000..34c1288 --- /dev/null +++ b/tests/test_server.py @@ -0,0 +1,105 @@ +import subprocess +import time +import json +import unittest +from unittest import mock + +import numpy as np +import evaluate +from whisper_live.server import TranscriptionServer +from whisper_live.client import TranscriptionClient +from whisper.normalizers import EnglishTextNormalizer + + +class TestTranscriptionServerInitialization(unittest.TestCase): + def test_initialization(self): + server = TranscriptionServer() + self.assertEqual(server.max_clients, 4) + self.assertEqual(server.max_connection_time, 600) + self.assertDictEqual(server.clients, {}) + self.assertDictEqual(server.websockets, {}) + self.assertDictEqual(server.clients_start_time, {}) + + +class TestGetWaitTime(unittest.TestCase): + def setUp(self): + self.server = TranscriptionServer() + self.server.clients_start_time = { + 'client1': time.time() - 120, + 'client2': time.time() - 300 + } + self.server.max_connection_time = 600 + + def test_get_wait_time(self): + expected_wait_time = (600 - (time.time() - self.server.clients_start_time['client2'])) / 60 + print(self.server.get_wait_time(), expected_wait_time) + self.assertAlmostEqual(self.server.get_wait_time(), expected_wait_time, places=2) + + +class TestServerConnection(unittest.TestCase): + def setUp(self): + self.server = TranscriptionServer() + + @mock.patch('websockets.WebSocketCommonProtocol') + def test_connection(self, mock_websocket): + mock_websocket.recv.return_value = json.dumps({ + 'uid': 'test_client', + 'language': 'en', + 'task': 'transcribe', + 'model': 'tiny.en' + }) + self.server.recv_audio(mock_websocket, "faster_whisper") + + + @mock.patch('websockets.WebSocketCommonProtocol') + def test_recv_audio_exception_handling(self, mock_websocket): + mock_websocket.recv.side_effect = [json.dumps({ + 'uid': 'test_client', + 'language': 'en', + 'task': 'transcribe', + 'model': 'tiny.en' + }), np.array([1, 2, 3]).tobytes()] + + with self.assertLogs(level="ERROR"): + self.server.recv_audio(mock_websocket, "faster_whisper") + + self.assertNotIn(mock_websocket, self.server.clients) + + +class TestServerInferenceAccuracy(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.server_process = subprocess.Popen(["python", "run_server.py"]) # Adjust the command as needed + time.sleep(2) + + @classmethod + def tearDownClass(cls): + cls.server_process.terminate() + cls.server_process.wait() + + @mock.patch('pyaudio.PyAudio') + def setUp(self, mock_pyaudio): + self.mock_pyaudio = mock_pyaudio.return_value + self.mock_stream = mock.MagicMock() + self.mock_pyaudio.open.return_value = self.mock_stream + self.metric = evaluate.load("wer") + self.normalizer = EnglishTextNormalizer() + self.client = TranscriptionClient( + "localhost", "9090", model="base.en", lang="en", + ) + + def test_inference(self): + gt = "And so my fellow Americans, ask not, what your country can do for you. Ask what you can do for your country!" + self.client("assets/jfk.flac") + with open("output.srt", "r") as f: + lines = f.readlines() + prediction = " ".join([l.strip() for l in lines[2::4]]) + prediction_normalized = self.normalizer(prediction) + gt_normalized = self.normalizer(gt) + + # calculate WER + wer = self.metric.compute( + predictions=[prediction_normalized], + references=[gt_normalized] + ) + self.assertLess(wer, 0.05)