102 lines
3.5 KiB
Python
102 lines
3.5 KiB
Python
import tempfile
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from aioresponses import CallbackResult, aioresponses
|
|
|
|
from beaver_gateway.transcribe import Audio, OpenAITranscriber, transcribe
|
|
|
|
URL = "http://whisper.local/v1/audio/transcriptions"
|
|
|
|
|
|
@pytest.fixture
|
|
def audio() -> Audio:
|
|
path = Path(tempfile.mkdtemp(prefix="beaver-stt-")) / "u1.ogg"
|
|
path.write_bytes(b"OggS")
|
|
return Audio(path, "audio/ogg", 3)
|
|
|
|
|
|
def saying(text: str | None, calls: list[str] | None = None):
|
|
async def transcriber(audio: Audio) -> str | None:
|
|
if calls is not None:
|
|
calls.append(text or "")
|
|
return text
|
|
|
|
return transcriber
|
|
|
|
|
|
async def broken(audio: Audio) -> str | None:
|
|
msg = "provider down"
|
|
raise ConnectionError(msg)
|
|
|
|
|
|
async def test_the_first_non_empty_text_wins(audio: Audio) -> None:
|
|
calls: list[str] = []
|
|
chain = [saying(None, calls), saying(" \n", calls), saying(" привет ", calls)]
|
|
chain.append(saying("never", calls))
|
|
assert await transcribe(chain, audio) == "привет"
|
|
assert calls == ["", " \n", " привет "]
|
|
|
|
|
|
async def test_a_raising_provider_is_skipped(audio: Audio, caplog) -> None:
|
|
assert await transcribe([broken, saying("ok")], audio) == "ok"
|
|
(record,) = [r for r in caplog.records if "provider down" in r.getMessage()]
|
|
assert record.levelname == "WARNING" and record.exc_info is None
|
|
|
|
|
|
async def test_nothing_comes_back_as_none(audio: Audio) -> None:
|
|
assert await transcribe([], audio) is None
|
|
assert await transcribe([broken, saying("")], audio) is None
|
|
|
|
|
|
async def test_openai_transcriber_posts_the_file_and_reads_the_text(
|
|
audio: Audio,
|
|
) -> None:
|
|
seen: dict[str, Any] = {}
|
|
|
|
def answer(url, **kwargs: Any) -> CallbackResult:
|
|
fields = kwargs["data"]._fields
|
|
seen["headers"] = kwargs["headers"]
|
|
seen["fields"] = {f[0]["name"]: f[2] for f in fields}
|
|
seen["file"] = next(f for f in fields if f[0]["name"] == "file")
|
|
return CallbackResult(payload={"text": " привет "})
|
|
|
|
stt = OpenAITranscriber(
|
|
base_url="http://whisper.local/v1/", api_key="sk", model="m", language="ru"
|
|
)
|
|
with aioresponses() as http:
|
|
http.post(URL, callback=answer)
|
|
assert await stt(audio) == " привет "
|
|
assert seen["headers"] == {"Authorization": "Bearer sk"}
|
|
assert seen["fields"]["model"] == "m"
|
|
assert seen["fields"]["language"] == "ru"
|
|
assert seen["fields"]["response_format"] == "json"
|
|
options, headers, _ = seen["file"]
|
|
assert options["filename"] == "u1.ogg"
|
|
assert headers["Content-Type"] == "audio/ogg"
|
|
|
|
|
|
async def test_openai_transcriber_without_a_key_sends_no_auth(audio: Audio) -> None:
|
|
seen: dict[str, Any] = {}
|
|
|
|
def answer(url, **kwargs: Any) -> CallbackResult:
|
|
seen["headers"] = kwargs["headers"]
|
|
seen["fields"] = [f[0]["name"] for f in kwargs["data"]._fields]
|
|
return CallbackResult(payload={"text": ""})
|
|
|
|
with aioresponses() as http:
|
|
http.post(URL, callback=answer)
|
|
assert (
|
|
await OpenAITranscriber(base_url="http://whisper.local/v1")(audio) is None
|
|
)
|
|
assert seen["headers"] == {}
|
|
assert "language" not in seen["fields"]
|
|
|
|
|
|
async def test_openai_transcriber_error_reaches_the_chain(audio: Audio) -> None:
|
|
stt = OpenAITranscriber(base_url="http://whisper.local/v1")
|
|
with aioresponses() as http:
|
|
http.post(URL, status=401, payload={"error": "bad key"})
|
|
assert await transcribe([stt, saying("fallback")], audio) == "fallback"
|