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