Add unit tests for BatchInferenceWorker

tests/test_batch_inference.py

@@ -0,0 +1,163 @@
+import time
+import unittest
+from unittest import mock
+from unittest.mock import MagicMock
+
+import numpy as np
+
+from whisper_live.batch_inference import BatchInferenceWorker, BatchRequest
+
+
+class TestBatchInferenceWorker(unittest.TestCase):
+    def setUp(self):
+        self.mock_transcriber = MagicMock()
+        self.worker = BatchInferenceWorker(
+            transcriber=self.mock_transcriber,
+            max_batch_size=8,
+            batch_window_ms=200,
+        )
+        self.worker.start()
+
+    def tearDown(self):
+        self.worker.stop()
+
+    def _make_audio(self, duration_s=1.0):
+        return np.random.randn(int(16000 * duration_s)).astype(np.float32)
+
+    def test_single_request_uses_transcribe(self):
+        """Single request should fall back to transcriber.transcribe()."""
+        fake_segment = MagicMock()
+        fake_info = MagicMock()
+        self.mock_transcriber.transcribe.return_value = ([fake_segment], fake_info)
+
+        req = BatchRequest(audio=self._make_audio(), language="en", use_vad=False)
+        self.worker.submit(req)
+        req.future.wait(timeout=5)
+
+        self.assertTrue(req.future.is_set())
+        self.assertIsNone(req.error)
+        self.assertEqual(req.result, [fake_segment])
+        self.assertEqual(req.info, fake_info)
+        self.mock_transcriber.transcribe.assert_called_once()
+
+    @mock.patch('whisper_live.batch_inference.get_suppressed_tokens', return_value=[-1])
+    @mock.patch('whisper_live.batch_inference.Tokenizer')
+    def test_multiple_requests_batched(self, mock_tokenizer_cls, mock_suppress):
+        """Multiple concurrent requests should go through the batched GPU path."""
+        # Mock tokenizer
+        mock_tok = MagicMock()
+        mock_tok.decode.return_value = "hello world"
+        mock_tokenizer_cls.return_value = mock_tok
+
+        # Mock feature extractor
+        self.mock_transcriber.feature_extractor.return_value = np.zeros(
+            (80, 3000), dtype=np.float32
+        )
+        self.mock_transcriber.feature_extractor.sampling_rate = 16000
+
+        # Mock encode
+        self.mock_transcriber.encode.return_value = np.zeros(
+            (3, 1500, 512), dtype=np.float32
+        )
+
+        # Mock model.generate — one result per item
+        gen_result = MagicMock()
+        gen_result.sequences_ids = [[50257, 50362, 1234, 50256]]
+        gen_result.scores = [np.float32(-1.0)]
+        gen_result.no_speech_prob = 0.1
+        self.mock_transcriber.model.generate.return_value = [gen_result] * 3
+
+        # Mock remaining model attributes
+        self.mock_transcriber.model.is_multilingual = False
+        self.mock_transcriber.max_length = 448
+        self.mock_transcriber.frames_per_second = 50
+        self.mock_transcriber.get_prompt.return_value = [50258]
+        self.mock_transcriber._split_segments_by_timestamps.return_value = (
+            [{"start": 0.0, "end": 1.0, "tokens": [1234], "seek": 0}],
+            None,
+            None,
+        )
+
+        requests = [
+            BatchRequest(audio=self._make_audio(), language="en", use_vad=False)
+            for _ in range(3)
+        ]
+        for req in requests:
+            self.worker.submit(req)
+        for req in requests:
+            req.future.wait(timeout=5)
+
+        for req in requests:
+            self.assertTrue(req.future.is_set())
+            self.assertIsNone(req.error)
+            self.assertIsNotNone(req.result)
+
+        # Verify the batched encode path was used (not transcribe)
+        self.mock_transcriber.encode.assert_called()
+        self.mock_transcriber.transcribe.assert_not_called()
+
+    def test_error_propagation(self):
+        """Transcriber errors should propagate to the request without crashing the worker."""
+        self.mock_transcriber.transcribe.side_effect = RuntimeError("GPU OOM")
+
+        req = BatchRequest(audio=self._make_audio(), language="en", use_vad=False)
+        self.worker.submit(req)
+        req.future.wait(timeout=5)
+
+        self.assertTrue(req.future.is_set())
+        self.assertIsInstance(req.error, RuntimeError)
+        self.assertIn("GPU OOM", str(req.error))
+
+        # Worker should still be alive — submit another request
+        self.mock_transcriber.transcribe.side_effect = None
+        self.mock_transcriber.transcribe.return_value = ([MagicMock()], MagicMock())
+
+        req2 = BatchRequest(audio=self._make_audio(), language="en", use_vad=False)
+        self.worker.submit(req2)
+        req2.future.wait(timeout=5)
+
+        self.assertIsNone(req2.error)
+        self.assertIsNotNone(req2.result)
+
+    def test_worker_stop(self):
+        """Worker thread should exit cleanly when stop() is called."""
+        self.assertTrue(self.worker._thread.is_alive())
+        self.worker.stop()
+        self.assertFalse(self.worker._thread.is_alive())
+
+    def test_batch_respects_max_size(self):
+        """Batches should not exceed max_batch_size."""
+        self.worker.stop()  # Stop the default worker
+
+        observed_batch_sizes = []
+        original_process = BatchInferenceWorker._process_batch
+
+        def tracking_process(self_inner, batch):
+            observed_batch_sizes.append(len(batch))
+            original_process(self_inner, batch)
+
+        self.worker = BatchInferenceWorker(
+            transcriber=self.mock_transcriber,
+            max_batch_size=2,
+            batch_window_ms=100,
+        )
+
+        self.mock_transcriber.transcribe.return_value = ([MagicMock()], MagicMock())
+
+        with mock.patch.object(
+            BatchInferenceWorker, '_process_batch', tracking_process
+        ):
+            self.worker.start()
+
+            requests = [
+                BatchRequest(audio=self._make_audio(), language="en", use_vad=False)
+                for _ in range(4)
+            ]
+            for req in requests:
+                self.worker.submit(req)
+            for req in requests:
+                req.future.wait(timeout=5)
+
+        for size in observed_batch_sizes:
+            self.assertLessEqual(size, 2)
+        self.assertTrue(all(req.future.is_set() for req in requests))