mistral-vibe/tests/transcribe/test_transcribe_client.py

from __future__ import annotations

from collections.abc import AsyncIterator
from unittest.mock import MagicMock

import pytest

from vibe.core.config import TranscribeModelConfig, TranscribeProviderConfig
from vibe.core.transcribe import (
    MistralTranscribeClient,
    TranscribeDone,
    TranscribeError,
    TranscribeSessionCreated,
    TranscribeTextDelta,
)


def _make_provider() -> TranscribeProviderConfig:
    return TranscribeProviderConfig(
        name="mistral",
        api_base="https://api.mistral.ai",
        api_key_env_var="MISTRAL_API_KEY",
    )


def _make_model() -> TranscribeModelConfig:
    return TranscribeModelConfig(
        name="mistral-small-transcribe",
        alias="foo",
        provider="mistral",
        encoding="pcm_s16le",
        sample_rate=16_000,
        target_streaming_delay_ms=200,
    )


async def _empty_audio_stream() -> AsyncIterator[bytes]:
    return
    yield


def _make_sdk_session_created(request_id: str = "test-request-id") -> MagicMock:
    from mistralai.client.models import (
        RealtimeTranscriptionSession,
        RealtimeTranscriptionSessionCreated,
    )

    session = MagicMock(spec=RealtimeTranscriptionSession)
    session.request_id = request_id
    mock = MagicMock(spec=RealtimeTranscriptionSessionCreated)
    mock.session = session
    return mock


def _make_sdk_text_delta(text: str) -> MagicMock:
    from mistralai.client.models import TranscriptionStreamTextDelta

    m = MagicMock(spec=TranscriptionStreamTextDelta)
    m.text = text
    return m


def _make_sdk_done(text: str) -> MagicMock:
    from mistralai.client.models import TranscriptionStreamDone

    m = MagicMock(spec=TranscriptionStreamDone)
    m.text = text
    return m


def _make_sdk_error(message: str) -> MagicMock:
    from mistralai.client.models import RealtimeTranscriptionError

    m = MagicMock(spec=RealtimeTranscriptionError)
    m.error = MagicMock()
    m.error.message = message
    return m


def _make_sdk_unknown() -> MagicMock:
    from mistralai.extra.realtime import UnknownRealtimeEvent

    return MagicMock(spec=UnknownRealtimeEvent)


async def _collect(client: MistralTranscribeClient) -> list[object]:
    events: list[object] = []
    async for event in client.transcribe(_empty_audio_stream()):
        events.append(event)
    return events


def _patch_sdk(sdk_events: list[object]) -> MagicMock:
    async def _fake_stream(**_kwargs: object) -> AsyncIterator[object]:
        for e in sdk_events:
            yield e

    mock_client = MagicMock()
    mock_client.audio.realtime.transcribe_stream = _fake_stream
    return mock_client


class TestEventMapping:
    @pytest.mark.asyncio
    async def test_session_created(self) -> None:
        mock_client = _patch_sdk([_make_sdk_session_created()])
        client = MistralTranscribeClient(_make_provider(), _make_model())
        client._client = mock_client

        events = await _collect(client)

        assert len(events) == 1
        assert isinstance(events[0], TranscribeSessionCreated)
        assert events[0].request_id == "test-request-id"

    @pytest.mark.asyncio
    async def test_text_delta(self) -> None:
        mock_client = _patch_sdk([_make_sdk_text_delta("hello")])
        client = MistralTranscribeClient(_make_provider(), _make_model())
        client._client = mock_client

        events = await _collect(client)

        assert len(events) == 1
        assert isinstance(events[0], TranscribeTextDelta)
        assert events[0].text == "hello"

    @pytest.mark.asyncio
    async def test_done(self) -> None:
        mock_client = _patch_sdk([_make_sdk_done("full text")])
        client = MistralTranscribeClient(_make_provider(), _make_model())
        client._client = mock_client

        events = await _collect(client)

        assert len(events) == 1
        assert isinstance(events[0], TranscribeDone)

    @pytest.mark.asyncio
    async def test_error(self) -> None:
        mock_client = _patch_sdk([_make_sdk_error("something broke")])
        client = MistralTranscribeClient(_make_provider(), _make_model())
        client._client = mock_client

        events = await _collect(client)

        assert len(events) == 1
        assert isinstance(events[0], TranscribeError)
        assert events[0].message == "something broke"

    @pytest.mark.asyncio
    async def test_unknown_event_is_skipped(self) -> None:
        mock_client = _patch_sdk([
            _make_sdk_session_created(),
            _make_sdk_unknown(),
            _make_sdk_text_delta("hi"),
        ])
        client = MistralTranscribeClient(_make_provider(), _make_model())
        client._client = mock_client

        events = await _collect(client)

        assert len(events) == 2
        assert isinstance(events[0], TranscribeSessionCreated)
        assert isinstance(events[1], TranscribeTextDelta)