satoko-chatbot/test_voice_stt.py

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

import httpx
import discord
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()


class CommandTests(unittest.IsolatedAsyncioTestCase):
    async def test_join_replies_in_command_channel_and_quit_closes_sink(self):
        # Charger uniquement la fonction : aucun token ni connexion Discord.
        tree = ast.parse(Path(__file__).with_name('chatbot.py').read_text(encoding='utf-8'))
        function = next(node for node in tree.body if isinstance(node, ast.AsyncFunctionDef)
                        and node.name == 'handle_voice_command')
        http = AsyncMock()
        http.__aenter__.return_value.get.return_value = Mock()
        sink = SimpleNamespace(close=AsyncMock())
        make_sink = Mock(return_value=sink)
        voice = SimpleNamespace(sink=sink, listen=Mock(), disconnect=AsyncMock())
        channel = SimpleNamespace(connect=AsyncMock(return_value=voice))
        voice.channel = channel
        message = SimpleNamespace(
            content='!join', guild=SimpleNamespace(voice_client=None),
            author=SimpleNamespace(voice=SimpleNamespace(channel=channel)),
            channel=SimpleNamespace(send=AsyncMock(), typing=AsyncMock()),
        )
        message.channel.typing = Mock(return_value=AsyncMock())
        call_api = AsyncMock(return_value=SimpleNamespace(
            choices=[SimpleNamespace(message=SimpleNamespace(content='Réponse'))]))
        namespace = dict(
            VOICE_ENABLED=True, STT_URL='http://stt/transcribe',
            voice_command_lock=asyncio.Lock(), httpx=SimpleNamespace(AsyncClient=Mock(return_value=http)),
            voice_recv=SimpleNamespace(VoiceRecvClient=object), STTSink=make_sink,
            call_openai_api=call_api, transform_emote=lambda text, output: text,
            filter_message=lambda text: text, split_message=lambda text: [text],
            discord=discord, logger=Mock(),
        )
        exec(compile(ast.Module(body=[function], type_ignores=[]), '<commands>', 'exec'), namespace)
        await namespace['handle_voice_command'](message)
        voice.listen.assert_called_once_with(sink)
        callback = make_sink.call_args.args[0]
        await callback(SimpleNamespace(display_name='Alice'), 'Bonjour')
        call_api.assert_awaited_once_with('Bonjour', 'Alice')
        self.assertEqual(message.channel.send.call_args.args[0], 'Réponse')
        message.content = '!quit'
        message.guild.voice_client = voice
        await namespace['handle_voice_command'](message)
        sink.close.assert_awaited_once()
        voice.disconnect.assert_awaited_once_with(force=True)


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