beaver-agent/tests/test_transcribe.py

import hashlib
import tempfile
from pathlib import Path

import pytest
from aiohttp import web
from aiohttp.test_utils import TestServer
from beaver_gateway.transcribe import Audio

from beaver_agent.transcribe import BeavergramTranscriber

OWNER = 777
VOICE = b"OggS voice"
DIGEST = hashlib.sha256(VOICE).hexdigest()


class Beavergram:
    """Поддельный beavergram: аккаунты, поиск по sha256 и расшифровка."""

    def __init__(self, *, empty_lookups=0, kind="voice", text=None, fail=None):
        self.empty_lookups = empty_lookups
        self.kind = kind
        self.text = text
        self.fail = fail
        self.lookups = []
        self.transcribes = []
        self.auth = set()

    def app(self):
        app = web.Application(middlewares=[self.gate])
        app.router.add_get("/api/accounts", self.accounts)
        app.router.add_get("/api/media/lookup", self.lookup)
        app.router.add_post("/api/media/transcribe", self.transcribe)
        return app

    @web.middleware
    async def gate(self, request, handler):
        self.auth.add(request.headers.get("Authorization"))
        if self.fail:
            return web.json_response({"detail": "nope"}, status=self.fail)
        return await handler(request)

    async def accounts(self, request):
        return web.json_response(
            [
                {"account_id": 1, "tg_user_id": 5, "is_bot": True},
                {"account_id": 2, "tg_user_id": OWNER, "is_bot": False},
            ]
        )

    async def lookup(self, request):
        self.lookups.append(dict(request.query))
        if len(self.lookups) <= self.empty_lookups:
            return web.json_response([])
        return web.json_response(
            [
                {
                    "account_id": 2,
                    "chat_id": 42,
                    "message_id": 9,
                    "kind": self.kind,
                    "extracted_text": self.text,
                }
            ]
        )

    async def transcribe(self, request):
        self.transcribes.append(await request.json())
        return web.json_response({"status": "done", "text": "привет из премиума"})


@pytest.fixture
def voice():
    path = Path(tempfile.mkdtemp(prefix="beaver-bg-")) / "1-vu1.ogg"
    path.write_bytes(VOICE)
    return path


async def serve(fake):
    server = TestServer(fake.app())
    await server.start_server()
    return server


def transcriber(server):
    return BeavergramTranscriber(
        origin=str(server.make_url("")).rstrip("/"),
        token="t0k",
        tg_user_id=OWNER,
        find_for=2.0,
        find_every=0.01,
    )


async def test_finds_the_voice_after_a_retry_and_asks_for_a_transcript(voice):
    fake = Beavergram(empty_lookups=2)
    server = await serve(fake)
    try:
        stt = transcriber(server)
        assert await stt(Audio(voice, "audio/ogg", 3)) == "привет из премиума"
    finally:
        await server.close()
    assert len(fake.lookups) == 3
    assert fake.lookups[-1] == {"sha256": DIGEST, "account_id": "2"}
    assert fake.transcribes == [
        {"account_id": 2, "chat_id": 42, "message_id": 9, "wait": 60}
    ]
    assert fake.auth == {"Bearer t0k"}


async def test_an_existing_transcript_comes_back_without_a_transcribe(voice):
    fake = Beavergram(kind="video_note", text="уже расшифровано")
    server = await serve(fake)
    try:
        assert (
            await transcriber(server)(Audio(voice, "video/mp4")) == "уже расшифровано"
        )
    finally:
        await server.close()
    assert fake.transcribes == []


async def test_music_and_other_files_are_not_touched(voice):
    fake = Beavergram()
    server = await serve(fake)
    try:
        stt = transcriber(server)
        for media_type in ("audio/mpeg", "audio/mp4", None):
            assert await stt(Audio(voice, media_type)) is None
        assert fake.lookups == []
        fake.kind = "audio"
        assert await stt(Audio(voice, "audio/ogg")) is None
    finally:
        await server.close()
    assert len(fake.lookups) == 1 and fake.transcribes == []


async def test_nothing_found_in_time_is_none(voice):
    fake = Beavergram(empty_lookups=1000)
    server = await serve(fake)
    try:
        stt = transcriber(server)
        stt.find_for = 0.05
        assert await stt(Audio(voice, "audio/ogg")) is None
    finally:
        await server.close()
    assert len(fake.lookups) >= 2 and fake.transcribes == []


@pytest.mark.parametrize("status", [401, 500])
async def test_http_errors_are_none_not_raised(voice, status):
    fake = Beavergram(fail=status)
    server = await serve(fake)
    try:
        assert await transcriber(server)(Audio(voice, "audio/ogg")) is None
    finally:
        await server.close()
    assert fake.auth == {"Bearer t0k"}


def test_from_env_needs_the_token_and_the_owner():
    url = "https://tg.example.com/mcp?token=" + "a" * 64
    stt = BeavergramTranscriber.from_env(
        {"BEAVERGRAM_MCP": url, "TELEGRAM_USER_ID": "777"}
    )
    assert stt is not None
    assert (stt.origin, stt.token, stt.tg_user_id) == (
        "https://tg.example.com",
        "a" * 64,
        777,
    )
    assert BeavergramTranscriber.from_env({"TELEGRAM_USER_ID": "777"}) is None
    assert (
        BeavergramTranscriber.from_env(
            {"BEAVERGRAM_MCP": "https://tg.example.com/mcp", "TELEGRAM_USER_ID": "777"}
        )
        is None
    )
    assert BeavergramTranscriber.from_env({"BEAVERGRAM_MCP": url}) is None