banban/talkingq-url/tests/test_fun_asr.py

"""Test real adapter I/O against a local DashScope-protocol WebSocket server."""
import asyncio
from types import SimpleNamespace

import pytest
import pytest_asyncio
from aiohttp import web, WSMsgType

from config import settings
from implementations.aliyun_asr import AliyunASR, ASRError
from services.audio_session import AudioSession
from services.interrupt_handler import interrupt_handler
from handlers.audio_session_handler import handle_audio_packet
from handlers.service_connection_handler import send_audio_task, end_asr_session


@pytest_asyncio.fixture(params=['qwen-audio-3.1-asr-flash-streaming', 'fun-asr-realtime'])
async def provider(monkeypatch, unused_tcp_port, request):
    state = SimpleNamespace(
        start_error=None, finish_error=None, start_delay=0, finish_delay=0,
        no_finish=False, early_close=False, events=[], chunks=[], commands=[],
        authorized=False, finish_seen=asyncio.Event(), started=asyncio.Event(), clients=[],
        model=request.param,
    )

    async def handler(request):
        state.authorized = request.headers.get('Authorization') == 'Bearer local-test-key'
        ws = web.WebSocketResponse()
        await ws.prepare(request)
        task_id = None

        async def emit(kind, payload=None, **header):
            await ws.send_json({'header': {'event': kind, 'task_id': task_id, **header},
                                'payload': payload or {}})

        async for message in ws:
            if message.type == WSMsgType.BINARY:
                state.chunks.append(message.data)
                if state.early_close:
                    await ws.close()
            elif message.type == WSMsgType.TEXT:
                command = message.json()
                state.commands.append(command)
                task_id = command['header']['task_id']
                if command['header']['action'] == 'run-task':
                    state.started.set()
                    await asyncio.sleep(state.start_delay)
                    if ws.closed:
                        break
                    if state.start_error:
                        await emit('task-failed', error_code=state.start_error, error_message='test rejection')
                    else:
                        await emit('task-started')
                elif command['header']['action'] == 'finish-task':
                    state.finish_seen.set()
                    await asyncio.sleep(state.finish_delay)
                    if ws.closed:
                        break
                    if state.finish_error:
                        await emit('task-failed', error_code=state.finish_error, error_message='test failure')
                    else:
                        await emit('result-generated', {'output': {'sentence': sentence(99, 'wrong task')}},
                                   task_id='stale-task')
                        for item in state.events:
                            await emit('result-generated', {'output': {'sentence': item}})
                        if not state.no_finish:
                            await emit('task-finished')
        return ws

    app = web.Application()
    app.router.add_get('/inference', handler)
    runner = web.AppRunner(app)
    await runner.setup()
    await web.TCPSite(runner, '127.0.0.1', unused_tcp_port).start()
    for key, value in {
        'aliyun_api_key': 'local-test-key',
        'aliyun_asr_url': f'http://127.0.0.1:{unused_tcp_port}/inference',
        'aliyun_asr_model': state.model,
        'aliyun_vocabulary_id': 'old-gummy-vocabulary', 'aliyun_asr_vocabulary_id': '',
    }.items():
        monkeypatch.setattr(settings, key, value)

    def client():
        asr = AliyunASR()
        state.clients.append(asr)
        return asr

    state.client = client
    yield state
    for asr in state.clients:
        await asr.close()
        assert asr._client is None and asr._receiver is None
    await runner.cleanup()


def sentence(number, text, final=True):
    return {'sentence_id': number, 'text': text, 'sentence_end': final}


@pytest.mark.asyncio
async def test_full_stream_keeps_all_sentences_and_tail(provider):
    provider.events = [sentence(1, '今天', False), sentence(1, '今天心情很好。'),
                       sentence(1, '今天心情很好。'), {'heartbeat': True, 'sentence_id': 0},
                       sentence(2, '请讲个', False), sentence(2, '请讲个故事。')]
    asr = provider.client()
    await asr.connect()
    audio = bytes(range(256)) * 29 + b'tail'
    for offset in range(0, len(audio), 257):
        await asr.send_audio(audio[offset:offset + 257])
    await asr.send_end()
    await asr.send_end()
    assert await asr.receive_results() == '今天心情很好。请讲个故事。'
    assert b''.join(provider.chunks) == audio
    assert provider.authorized
    request = provider.commands[0]['payload']
    assert request['model'] == provider.model and request['function'] == 'recognition'
    assert request['parameters']['format'] == 'mp3'
    assert request['parameters']['sample_rate'] == 16000
    assert 'vocabulary_id' not in request['parameters']
    assert sum(c['header']['action'] == 'finish-task' for c in provider.commands) == 1


@pytest.mark.asyncio
async def test_explicit_new_model_vocabulary(provider, monkeypatch):
    monkeypatch.setattr(settings, 'aliyun_asr_vocabulary_id', 'fun-asr-hotwords')
    asr = provider.client()
    await asr.connect()
    assert provider.commands[0]['payload']['parameters']['vocabulary_id'] == 'fun-asr-hotwords'


@pytest.mark.asyncio
async def test_waits_for_provider_task_started(provider):
    provider.start_delay = 0.05
    asr = provider.client()
    connecting = asyncio.create_task(asr.connect())
    await provider.started.wait()
    assert not asr.connected
    await connecting
    assert asr.connected


@pytest.mark.asyncio
@pytest.mark.parametrize('code', ['InvalidParameter', 'InvalidApiKey', 'Throttling'])
async def test_start_rejection_does_not_report_connected(provider, code):
    provider.start_error = code
    asr = provider.client()
    with pytest.raises(ASRError, match=code):
        await asr.connect()
    assert not asr.connected and asr._client is None and asr.transcript == ''


@pytest.mark.asyncio
async def test_connect_timeout_releases_resources(provider, monkeypatch):
    provider.start_delay = 0.1
    monkeypatch.setattr(settings, 'asr_connect_timeout_seconds', 0.02)
    asr = provider.client()
    with pytest.raises(ASRError, match='TimeoutError'):
        await asr.connect()
    assert asr._client is None


@pytest.mark.asyncio
async def test_final_after_old_half_second_timeout_is_collected(provider):
    provider.finish_delay = 0.6
    provider.events = [sentence(1, '请完整保留最后一句。')]
    asr = provider.client()
    await asr.connect()
    await asr.send_audio(b'short')
    await asr.send_end()
    assert await asr.receive_results() == '请完整保留最后一句。'
    assert provider.chunks == [b'short']


@pytest.mark.asyncio
async def test_final_timeout_rejects_partial_transcript(provider, monkeypatch):
    provider.no_finish = True
    provider.events = [sentence(1, '半句话', False)]
    monkeypatch.setattr(settings, 'asr_final_timeout_seconds', 0.03)
    asr = provider.client()
    await asr.connect()
    with pytest.raises(ASRError, match='TimeoutError'):
        await asr.send_end()
    with pytest.raises(ASRError):
        await asr.receive_results()
    assert asr.transcript == ''


@pytest.mark.asyncio
async def test_finish_without_final_sentence_is_rejected(provider):
    provider.events = [sentence(1, '未定稿', False)]
    asr = provider.client()
    await asr.connect()
    await asr.send_end()
    with pytest.raises(ASRError, match='不完整'):
        await asr.receive_results()


@pytest.mark.asyncio
async def test_provider_failure_during_finish(provider):
    provider.finish_error = 'InternalError'
    asr = provider.client()
    await asr.connect()
    with pytest.raises(ASRError, match='InternalError'):
        await asr.send_end()
    assert not asr.connected


@pytest.mark.asyncio
async def test_early_disconnect_is_not_success(provider):
    provider.early_close = True
    asr = provider.client()
    await asr.connect()
    await asr.send_audio(b'x' * 3200)
    await asyncio.wait_for(asr._finished.wait(), 1)
    with pytest.raises(ASRError, match='提前关闭'):
        await asr.receive_results()


@pytest.mark.asyncio
async def test_silence_has_empty_transcript(provider):
    asr = provider.client()
    await asr.connect()
    await asr.send_end()
    assert await asr.receive_results() == ''


@pytest.mark.asyncio
async def test_cancel_finish_releases_socket(provider):
    provider.no_finish = True
    asr = provider.client()
    await asr.connect()
    ending = asyncio.create_task(asr.send_end())
    await provider.finish_seen.wait()
    ending.cancel()
    with pytest.raises(asyncio.CancelledError):
        await ending
    assert asr._client is None and not asr.connected


@pytest.mark.asyncio
async def test_reconnect_does_not_reuse_previous_transcript(provider):
    asr = provider.client()
    provider.events = [sentence(1, '第一轮。')]
    await asr.connect()
    await asr.send_end()
    assert await asr.receive_results() == '第一轮。'
    old_task = asr.task_id
    await asr.close()
    provider.events = [sentence(1, '第二轮。')]
    await asr.connect()
    await asr.send_end()
    assert await asr.receive_results() == '第二轮。'
    assert old_task != asr.task_id


def audio_session(asr):
    session = AudioSession(asr_service=asr)
    session.device_id, session.session_id = 'local-test', 'round'
    session.send_queue = asyncio.Queue()
    session.send_task = asyncio.create_task(send_audio_task(session))
    return session


@pytest.mark.asyncio
async def test_end_drains_queued_audio_before_finish(provider):
    provider.events = [sentence(1, '完整识别。')]
    asr = provider.client()
    await asr.connect()
    session = audio_session(asr)
    # No yield to the consumer: this reproduced the previous cancellation loss.
    for part in [b'a' * 3200, b'b' * 3200, b'last']:
        await handle_audio_packet(session, part)
    await end_asr_session(session)
    assert b''.join(provider.chunks) == b'a' * 3200 + b'b' * 3200 + b'last'
    assert session.transcript == '完整识别。' and session.asr_error is None
    assert session.send_task is None and session.send_queue.empty()
    await asyncio.wait_for(session.send_queue.join(), 0.1)
    assert asr._client is None


@pytest.mark.asyncio
async def test_queue_timeout_closes_asr_without_partial_result(provider, monkeypatch):
    monkeypatch.setattr(settings, 'asr_queue_drain_timeout_seconds', 0.02)
    asr = provider.client()
    await asr.connect()
    session = audio_session(asr)

    async def stalled(data):
        await asyncio.Event().wait()

    monkeypatch.setattr(asr, 'send_audio', stalled)
    await handle_audio_packet(session, b'pending')
    await end_asr_session(session)
    assert 'TimeoutError' in session.asr_error
    assert session.transcript == '' and not provider.finish_seen.is_set()
    assert asr._client is None


@pytest.mark.asyncio
async def test_queue_overflow_does_not_return_incomplete_recording(provider):
    asr = provider.client()
    await asr.connect()
    session = audio_session(asr)
    for _ in range(101):
        await handle_audio_packet(session, b'x')
    await end_asr_session(session)
    assert '溢出' in session.asr_error and session.transcript == ''
    assert asr._client is None


@pytest.mark.asyncio
async def test_sender_failure_still_closes_connection(provider, monkeypatch):
    asr = provider.client()
    await asr.connect()
    session = audio_session(asr)

    async def fail(data):
        raise RuntimeError('broken send')

    monkeypatch.setattr(asr, 'send_audio', fail)
    await handle_audio_packet(session, b'audio')
    await end_asr_session(session)
    assert 'broken send' in session.asr_error and session.transcript == ''
    assert asr._client is None


@pytest.mark.asyncio
async def test_interruption_closes_asr_without_transcript(provider):
    asr = provider.client()
    await asr.connect()
    session = audio_session(asr)
    key = (session.device_id, session.session_id)
    await interrupt_handler.set_interrupt_state(key, True)
    try:
        await end_asr_session(session)
        assert '中断' in session.asr_error and session.transcript == ''
        assert asr._client is None
    finally:
        await interrupt_handler.remove_session(key)


@pytest.mark.asyncio
async def test_cancel_connect_releases_socket(provider):
    provider.start_delay = 0.1
    asr = provider.client()
    connecting = asyncio.create_task(asr.connect())
    await provider.started.wait()
    connecting.cancel()
    with pytest.raises(asyncio.CancelledError):
        await connecting
    assert asr._client is None and asr._receiver is None


@pytest.mark.asyncio
async def test_concurrent_cleanup_is_idempotent(provider):
    asr = provider.client()
    await asr.connect()
    await asyncio.gather(asr.close(), asr.close())
    assert asr._client is None and asr._receiver is None


@pytest.mark.asyncio
async def test_send_timeout_is_reported(provider, monkeypatch):
    monkeypatch.setattr(settings, 'asr_send_timeout_seconds', 0.02)
    asr = provider.client()
    await asr.connect()

    async def blocked(data):
        await asyncio.Event().wait()

    monkeypatch.setattr(asr._websocket, 'send_bytes', blocked)
    with pytest.raises(ASRError, match='TimeoutError'):
        await asr.send_audio(b'a' * 3200)
    assert not asr.connected


@pytest.mark.asyncio
async def test_active_interrupt_releases_asr_and_sender(provider):
    from handlers.service_connection_handler import prepare_asr_service
    asr = provider.client()
    session = AudioSession(asr_service=asr)
    session.device_id, session.session_id = 'local-test', 'active-interrupt'
    session.send_queue = asyncio.Queue()
    key = (session.device_id, session.session_id)
    try:
        await prepare_asr_service(session, *key)
        assert asr.connected and session.send_task
        await interrupt_handler.handle_interrupt(key)
        assert asr._client is None and asr._receiver is None
        assert session.send_task is None
    finally:
        await interrupt_handler.remove_session(key)


@pytest.mark.asyncio
async def test_start_failure_does_not_leave_waiting_sender(provider):
    from handlers.service_connection_handler import prepare_asr_service
    provider.start_error = 'InvalidParameter'
    asr = provider.client()
    session = AudioSession(asr_service=asr)
    session.device_id, session.session_id = 'local-test', 'failed-start'
    session.send_queue = asyncio.Queue()
    await prepare_asr_service(session, session.device_id, session.session_id)
    assert 'InvalidParameter' in session.asr_error
    assert session.send_task is None and asr._client is None


@pytest.mark.asyncio
async def test_device_end_handler_passes_full_text_to_response_stage(provider, monkeypatch):
    from unittest.mock import AsyncMock
    from handlers.audio_session_handler import handle_end_packet
    import handlers.transcription_handler as transcription_handler
    provider.events = [sentence(1, '第一句。'), sentence(2, '第二句。')]
    asr = provider.client()
    await asr.connect()
    session = audio_session(asr)
    handle_transcription = AsyncMock()
    monkeypatch.setattr(transcription_handler, 'handle_transcription', handle_transcription)
    await handle_audio_packet(session, b'last bytes')
    await handle_end_packet(session, (session.device_id, session.session_id))
    handle_transcription.assert_awaited_once_with(session, session.device_id, session.session_id)
    assert session.transcript == '第一句。第二句。'
    assert b''.join(provider.chunks) == b'last bytes'


@pytest.mark.asyncio
async def test_empty_recognition_uses_existing_prompt_fallback(provider, monkeypatch):
    from unittest.mock import AsyncMock
    from handlers.audio_session_handler import handle_end_packet
    import handlers.prompt_sound_handler as prompts
    from services.tts_error_manager import tts_error_manager
    asr = provider.client()
    await asr.connect()
    session = audio_session(asr)
    prompt = AsyncMock()
    monkeypatch.setattr(prompts, 'handle_prompt_sound_request', prompt)
    monkeypatch.setattr(tts_error_manager, 'end_tts_session', AsyncMock())
    await handle_end_packet(session, (session.device_id, session.session_id))
    prompt.assert_awaited_once_with(session.device_id, 'tts_error')
    assert session.transcript == '' and asr._client is None


@pytest.mark.asyncio
async def test_cancel_session_start_cleans_up_sender(provider):
    from handlers.service_connection_handler import prepare_asr_service
    provider.start_delay = 0.1
    asr = provider.client()
    session = AudioSession(asr_service=asr)
    session.device_id, session.session_id = 'local-test', 'cancel-start'
    session.send_queue = asyncio.Queue()
    task = asyncio.create_task(prepare_asr_service(session, session.device_id, session.session_id))
    await provider.started.wait()
    task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await task
    assert session.send_task is None and asr._client is None