satoko-chatbot/test_voice_stt.py

"""Tests hors Discord : python -m unittest test_voice_stt."""
import asyncio
import time
import unittest
import wave
from io import BytesIO
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock, patch

import httpx
from voice_stt import STTSink


class VoiceTests(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.reply = AsyncMock()
        self.sink = STTSink(self.reply, 'http://stt/transcribe')
        self.user = SimpleNamespace(id=1, bot=False)
        self.pcm = b'\x00\x10' * (48000 * 2 // 2)  # 0,5 s stéréo

    async def asyncTearDown(self):
        await self.sink.close()

    async def test_silence_and_separate_speakers(self):
        self.sink.worker_task.cancel()
        with self.sink.lock:
            self.sink.add_pcm(self.user, bytes(len(self.pcm)), time.monotonic())
            self.assertFalse(self.sink.buffers)
            self.sink.add_pcm(self.user, self.pcm, time.monotonic() - 2)
            self.sink.add_pcm(SimpleNamespace(id=2), self.pcm, time.monotonic())
        await asyncio.sleep(0.15)
        user, pcm = self.sink.queue.get_nowait()
        self.sink.queue.task_done()
        self.assertEqual(user.id, 1)
        self.assertEqual(pcm, self.pcm)
        self.assertIn(2, self.sink.buffers)

    async def test_audio_is_bounded_and_short_noise_ignored(self):
        self.sink.worker_task.cancel()
        with self.sink.lock:
            self.sink.add_pcm(self.user, self.pcm * 50, time.monotonic())
            self.assertEqual(len(self.sink.buffers[1][1]), 48000 * 4 * 20)
            self.sink.add_pcm(SimpleNamespace(id=2), self.pcm[:1000], time.monotonic() - 2)
        await asyncio.sleep(0.15)
        self.assertEqual(self.sink.queue.qsize(), 1)
        self.assertNotIn(2, self.sink.buffers)

    async def test_dave_decrypted_before_opus(self):
        session = SimpleNamespace(ready=True, decrypt=Mock(return_value=b'opus'))
        self.sink._voice_client = SimpleNamespace(_connection=SimpleNamespace(
            dave_protocol_version=1, dave_session=session))
        data = SimpleNamespace(opus=b'encrypted', packet=SimpleNamespace(is_silence=lambda: False))
        decoder = Mock()
        decoder.decode.return_value = self.pcm
        with patch('voice_stt.discord.opus.Decoder', return_value=decoder):
            self.sink.write(self.user, data)
        session.decrypt.assert_called_once()
        decoder.decode.assert_called_once_with(b'opus', fec=False)
        self.assertEqual(bytes(self.sink.buffers[1][1]), self.pcm)

    async def test_http_recovery_and_wav_format(self):
        await self.sink.close()
        post = AsyncMock(side_effect=[
            httpx.ConnectError('offline'),
            httpx.Response(200, json={'text': 'Bonjour'}, request=httpx.Request('POST', 'http://stt')),
        ])
        http = AsyncMock()
        http.__aenter__.return_value.post = post
        with patch('voice_stt.httpx.AsyncClient', return_value=http):
            self.sink = STTSink(self.reply, 'http://stt/transcribe')
            self.sink.queue.put_nowait((self.user, self.pcm))
            self.sink.queue.put_nowait((self.user, self.pcm))
            await asyncio.wait_for(self.sink.queue.join(), 2)
            self.reply.assert_awaited_once_with(self.user, 'Bonjour')
            payload = post.call_args.kwargs['files']['file'][1]
            with wave.open(BytesIO(payload)) as wav:
                self.assertEqual((wav.getnchannels(), wav.getsampwidth(), wav.getframerate()), (2, 2, 48000))
                self.assertEqual(wav.readframes(wav.getnframes()), self.pcm)

    async def test_close_discards_pending_audio(self):
        self.sink.queue.put_nowait((self.user, self.pcm))
        await self.sink.close()
        self.assertTrue(self.sink.closed)
        self.assertTrue(self.sink.queue.empty())
        self.assertTrue(self.sink.flush_task.done())
        self.assertTrue(self.sink.worker_task.done())
        self.reply.assert_not_awaited()


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