Fix unittest to exposed client manager args

Signed-off-by: makaveli10 <suryanvineet47@gmail.com>
This commit is contained in:
makaveli10
2024-10-28 17:01:06 +05:30
parent 617fda2864
commit 8b87a0562d
2 changed files with 15 additions and 15 deletions
+3 -1
View File
@@ -48,7 +48,9 @@ class TestClientCallbacks(BaseTestCase):
"language": self.client.language, "language": self.client.language,
"task": self.client.task, "task": self.client.task,
"model": self.client.model, "model": self.client.model,
"use_vad": True "use_vad": True,
"max_clients": 4,
"max_connection_time": 600,
}) })
self.client.on_open(self.mock_ws_app) self.client.on_open(self.mock_ws_app)
self.mock_ws_app.send.assert_called_with(expected_message) self.mock_ws_app.send.assert_called_with(expected_message)
+12 -14
View File
@@ -5,10 +5,10 @@ import unittest
from unittest import mock from unittest import mock
import numpy as np import numpy as np
import evaluate import jiwer
from websockets.exceptions import ConnectionClosed from websockets.exceptions import ConnectionClosed
from whisper_live.server import TranscriptionServer from whisper_live.server import TranscriptionServer, BackendType, ClientManager
from whisper_live.client import Client, TranscriptionClient, TranscriptionTeeClient from whisper_live.client import Client, TranscriptionClient, TranscriptionTeeClient
from whisper.normalizers import EnglishTextNormalizer from whisper.normalizers import EnglishTextNormalizer
@@ -16,6 +16,7 @@ from whisper.normalizers import EnglishTextNormalizer
class TestTranscriptionServerInitialization(unittest.TestCase): class TestTranscriptionServerInitialization(unittest.TestCase):
def test_initialization(self): def test_initialization(self):
server = TranscriptionServer() server = TranscriptionServer()
server.client_manager = ClientManager(max_clients=4, max_connection_time=600)
self.assertEqual(server.client_manager.max_clients, 4) self.assertEqual(server.client_manager.max_clients, 4)
self.assertEqual(server.client_manager.max_connection_time, 600) self.assertEqual(server.client_manager.max_connection_time, 600)
self.assertDictEqual(server.client_manager.clients, {}) self.assertDictEqual(server.client_manager.clients, {})
@@ -25,6 +26,7 @@ class TestTranscriptionServerInitialization(unittest.TestCase):
class TestGetWaitTime(unittest.TestCase): class TestGetWaitTime(unittest.TestCase):
def setUp(self): def setUp(self):
self.server = TranscriptionServer() self.server = TranscriptionServer()
self.server.client_manager = ClientManager(max_clients=4, max_connection_time=600)
self.server.client_manager.start_times = { self.server.client_manager.start_times = {
'client1': time.time() - 120, 'client1': time.time() - 120,
'client2': time.time() - 300 'client2': time.time() - 300
@@ -49,7 +51,7 @@ class TestServerConnection(unittest.TestCase):
'task': 'transcribe', 'task': 'transcribe',
'model': 'tiny.en' 'model': 'tiny.en'
}) })
self.server.recv_audio(mock_websocket, "faster_whisper") self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
@mock.patch('websockets.WebSocketCommonProtocol') @mock.patch('websockets.WebSocketCommonProtocol')
def test_recv_audio_exception_handling(self, mock_websocket): def test_recv_audio_exception_handling(self, mock_websocket):
@@ -61,7 +63,7 @@ class TestServerConnection(unittest.TestCase):
}), np.array([1, 2, 3]).tobytes()] }), np.array([1, 2, 3]).tobytes()]
with self.assertLogs(level="ERROR"): with self.assertLogs(level="ERROR"):
self.server.recv_audio(mock_websocket, "faster_whisper") self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
self.assertNotIn(mock_websocket, self.server.client_manager.clients) self.assertNotIn(mock_websocket, self.server.client_manager.clients)
@@ -82,7 +84,6 @@ class TestServerInferenceAccuracy(unittest.TestCase):
cls.server_process.wait() cls.server_process.wait()
def setUp(self): def setUp(self):
self.metric = evaluate.load("wer")
self.normalizer = EnglishTextNormalizer() self.normalizer = EnglishTextNormalizer()
def check_prediction(self, srt_path): def check_prediction(self, srt_path):
@@ -94,11 +95,8 @@ class TestServerInferenceAccuracy(unittest.TestCase):
gt_normalized = self.normalizer(gt) gt_normalized = self.normalizer(gt)
# calculate WER # calculate WER
wer = self.metric.compute( wer_score = jiwer.wer(gt_normalized, prediction_normalized)
predictions=[prediction_normalized], self.assertLess(wer_score, 0.05)
references=[gt_normalized]
)
self.assertLess(wer, 0.05)
def test_inference(self): def test_inference(self):
client = TranscriptionClient( client = TranscriptionClient(
@@ -124,10 +122,10 @@ class TestExceptionHandling(unittest.TestCase):
@mock.patch('websockets.WebSocketCommonProtocol') @mock.patch('websockets.WebSocketCommonProtocol')
def test_connection_closed_exception(self, mock_websocket): def test_connection_closed_exception(self, mock_websocket):
mock_websocket.recv.side_effect = ConnectionClosed(1001, "testing connection closed") mock_websocket.recv.side_effect = ConnectionClosed(1001, "testing connection closed", rcvd_then_sent=mock.Mock())
with self.assertLogs(level="INFO") as log: with self.assertLogs(level="INFO") as log:
self.server.recv_audio(mock_websocket, "faster_whisper") self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
self.assertTrue(any("Connection closed by client" in message for message in log.output)) self.assertTrue(any("Connection closed by client" in message for message in log.output))
@mock.patch('websockets.WebSocketCommonProtocol') @mock.patch('websockets.WebSocketCommonProtocol')
@@ -135,7 +133,7 @@ class TestExceptionHandling(unittest.TestCase):
mock_websocket.recv.return_value = "invalid json" mock_websocket.recv.return_value = "invalid json"
with self.assertLogs(level="ERROR") as log: with self.assertLogs(level="ERROR") as log:
self.server.recv_audio(mock_websocket, "faster_whisper") self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
self.assertTrue(any("Failed to decode JSON from client" in message for message in log.output)) self.assertTrue(any("Failed to decode JSON from client" in message for message in log.output))
@mock.patch('websockets.WebSocketCommonProtocol') @mock.patch('websockets.WebSocketCommonProtocol')
@@ -143,7 +141,7 @@ class TestExceptionHandling(unittest.TestCase):
mock_websocket.recv.side_effect = RuntimeError("Unexpected error") mock_websocket.recv.side_effect = RuntimeError("Unexpected error")
with self.assertLogs(level="ERROR") as log: with self.assertLogs(level="ERROR") as log:
self.server.recv_audio(mock_websocket, "faster_whisper") self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
for message in log.output: for message in log.output:
print(message) print(message)
print() print()