tests/test_media_queue.py - master - artyuhovs/AIFeedBackTrainingBot

tests/test_media_queue.py
340 строк12 KB
Codex
Добавить прозрачный медиаконвейер в версии 0.8.1
19 июл 2026, 20:24
19 июл 2026, 20:249c5c8db
from __future__ import annotations

import asyncio
from collections.abc import Callable
from datetime import UTC, datetime
from pathlib import Path
from types import SimpleNamespace

import pytest

from app.db.repositories import Repositories
from app.models.media import IncomingMedia
from app.models.transcription import PreparedAudio, TranscriptionResult
from app.services.media_queue import MediaProcessingQueue, StageQueuePosition
from app.services.pipeline import ProcessingPipeline


async def _wait_until(predicate: Callable[[], bool]) -> None:
    async def wait() -> None:
        while not predicate():
            await asyncio.sleep(0)

    await asyncio.wait_for(wait(), timeout=1)


def _position(value: StageQueuePosition) -> tuple[bool, int | None]:
    return value.active, value.position


def test_media_queue_reports_waiting_positions_and_starts_in_fifo_order() -> None:
    async def run() -> tuple[list[str], dict[str, list[tuple[bool, int | None]]]]:
        queue = MediaProcessingQueue()
        entered: list[str] = []
        positions: dict[str, list[tuple[bool, int | None]]] = {name: [] for name in ("first", "second", "third")}
        releases = {name: asyncio.Event() for name in positions}

        async def worker(name: str) -> None:
            async def report(position: StageQueuePosition) -> None:
                positions[name].append(_position(position))

            async with queue.acquire(report):
                entered.append(name)
                await releases[name].wait()

        first = asyncio.create_task(worker("first"))
        await _wait_until(lambda: entered == ["first"])
        second = asyncio.create_task(worker("second"))
        await _wait_until(lambda: positions["second"] == [(False, 1)])
        third = asyncio.create_task(worker("third"))
        await _wait_until(lambda: positions["third"] == [(False, 2)])

        releases["first"].set()
        await _wait_until(
            lambda: (
                entered == ["first", "second"]
                and positions["second"][-1] == (True, None)
                and positions["third"][-1] == (False, 1)
            )
        )
        releases["second"].set()
        await _wait_until(lambda: entered == ["first", "second", "third"] and positions["third"][-1] == (True, None))
        releases["third"].set()
        await asyncio.gather(first, second, third)
        return entered, positions

    entered, positions = asyncio.run(run())

    assert entered == ["first", "second", "third"]
    assert positions == {
        "first": [(True, None)],
        "second": [(False, 1), (True, None)],
        "third": [(False, 2), (False, 1), (True, None)],
    }


def test_media_queue_allows_configured_parallelism_and_keeps_waiters_fifo() -> None:
    async def run() -> tuple[list[str], int, list[tuple[bool, int | None]]]:
        queue = MediaProcessingQueue(concurrency=2)
        entered: list[str] = []
        active = 0
        max_active = 0
        releases = {name: asyncio.Event() for name in ("first", "second", "third")}
        third_positions: list[tuple[bool, int | None]] = []

        async def worker(name: str) -> None:
            nonlocal active, max_active

            async def report(position: StageQueuePosition) -> None:
                if name == "third":
                    third_positions.append(_position(position))

            async with queue.acquire(report):
                entered.append(name)
                active += 1
                max_active = max(max_active, active)
                try:
                    await releases[name].wait()
                finally:
                    active -= 1

        first = asyncio.create_task(worker("first"))
        second = asyncio.create_task(worker("second"))
        await _wait_until(lambda: len(entered) == 2)
        third = asyncio.create_task(worker("third"))
        await _wait_until(lambda: third_positions == [(False, 1)])

        releases["first"].set()
        await _wait_until(lambda: "third" in entered and third_positions[-1] == (True, None))
        releases["second"].set()
        releases["third"].set()
        await asyncio.gather(first, second, third)
        return entered, max_active, third_positions

    entered, max_active, third_positions = asyncio.run(run())

    assert entered[:2] == ["first", "second"]
    assert entered[2] == "third"
    assert max_active == 2
    assert third_positions == [(False, 1), (True, None)]


def test_cancelling_a_waiter_updates_the_users_behind_it() -> None:
    async def run() -> tuple[list[str], list[tuple[bool, int | None]]]:
        queue = MediaProcessingQueue()
        entered: list[str] = []
        release_first = asyncio.Event()
        release_third = asyncio.Event()
        third_positions: list[tuple[bool, int | None]] = []

        async def worker(name: str, release: asyncio.Event) -> None:
            async def report(position: StageQueuePosition) -> None:
                if name == "third":
                    third_positions.append(_position(position))

            async with queue.acquire(report):
                entered.append(name)
                await release.wait()

        first = asyncio.create_task(worker("first", release_first))
        await _wait_until(lambda: entered == ["first"])
        second = asyncio.create_task(worker("second", asyncio.Event()))
        await asyncio.sleep(0)
        third = asyncio.create_task(worker("third", release_third))
        await _wait_until(lambda: third_positions == [(False, 2)])

        second.cancel()
        with pytest.raises(asyncio.CancelledError):
            await second
        await _wait_until(lambda: third_positions[-1] == (False, 1))

        release_first.set()
        await _wait_until(lambda: entered == ["first", "third"])
        release_third.set()
        await asyncio.gather(first, third)
        return entered, third_positions

    entered, third_positions = asyncio.run(run())

    assert entered == ["first", "third"]
    assert third_positions == [(False, 2), (False, 1), (True, None)]


def test_active_failure_releases_the_next_media_item() -> None:
    async def run() -> tuple[list[str], list[tuple[bool, int | None]]]:
        queue = MediaProcessingQueue()
        entered: list[str] = []
        second_positions: list[tuple[bool, int | None]] = []
        fail_first = asyncio.Event()

        async def first_worker() -> None:
            async with queue.acquire():
                entered.append("first")
                await fail_first.wait()
                raise RuntimeError("processing failed")

        async def second_worker() -> None:
            async def report(position: StageQueuePosition) -> None:
                second_positions.append(_position(position))

            async with queue.acquire(report):
                entered.append("second")

        first = asyncio.create_task(first_worker())
        await _wait_until(lambda: entered == ["first"])
        second = asyncio.create_task(second_worker())
        await _wait_until(lambda: second_positions == [(False, 1)])
        fail_first.set()
        with pytest.raises(RuntimeError):
            await first
        await second
        return entered, second_positions

    entered, second_positions = asyncio.run(run())

    assert entered == ["first", "second"]
    assert second_positions == [(False, 1), (True, None)]


def test_position_callback_failure_does_not_block_processing() -> None:
    async def run() -> bool:
        queue = MediaProcessingQueue()

        async def failing_callback(_position: StageQueuePosition) -> None:
            raise RuntimeError("telegram unavailable")

        async with queue.acquire(failing_callback):
            return True

    assert asyncio.run(run()) is True


def test_pipeline_overlaps_download_audio_and_asr_stages(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    repository = Repositories(tmp_path / "bot.sqlite3")
    repository.schema.init()
    for user_id in (41, 42, 43):
        repository.identities.upsert_user(user_id, f"user-{user_id}", "User", "")
        repository.input_intents.start_free_training(user_id)

    settings = SimpleNamespace(
        workdir=tmp_path / "files",
        debug_save_raw=False,
        media_pipeline_concurrency=8,
        media_download_concurrency=4,
        media_audio_concurrency=2,
        media_asr_concurrency=1,
    )
    all_downloads_started = asyncio.Event()
    asr_started = asyncio.Event()
    second_audio_started = asyncio.Event()
    overlap_observed = asyncio.Event()
    stage_active = {"download": 0, "audio": 0, "asr": 0}
    stage_max = {"download": 0, "audio": 0, "asr": 0}
    overlap_snapshot: dict[str, int] = {}

    def enter(stage: str) -> None:
        stage_active[stage] += 1
        stage_max[stage] = max(stage_max[stage], stage_active[stage])

    def leave(stage: str) -> None:
        stage_active[stage] -= 1

    async def fake_download_media(
        _bot: object,
        _settings: object,
        meta: object,
        record_id: int,
    ) -> IncomingMedia:
        enter("download")
        try:
            if stage_active["download"] == 3:
                all_downloads_started.set()
            await all_downloads_started.wait()
            message_id = int(getattr(meta, "telegram_message_id"))
            if message_id == 2:
                await asr_started.wait()
            elif message_id == 3:
                await second_audio_started.wait()
                overlap_snapshot.update(stage_active)
                overlap_observed.set()
            path = settings.workdir / str(record_id) / f"original-{message_id}.mp4"
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_bytes(b"video")
            return IncomingMedia(**meta.model_dump(), local_path=path)
        finally:
            leave("download")

    class ControlledAudioConverter:
        async def prepare(self, input_path: Path, output_dir: Path) -> PreparedAudio:
            enter("audio")
            try:
                message_id = int(input_path.stem.rsplit("-", 1)[1])
                if message_id == 2:
                    second_audio_started.set()
                    await overlap_observed.wait()
                pcm_path = output_dir / f"audio-{message_id}.pcm"
                wav_path = output_dir / f"audio-{message_id}.wav"
                pcm_path.write_bytes(b"pcm")
                wav_path.write_bytes(b"wav")
                return PreparedAudio(path=pcm_path, wav_path=wav_path, duration_sec=2, size_bytes=3)
            finally:
                leave("audio")

    class ControlledASR:
        async def transcribe(self, prepared: PreparedAudio) -> TranscriptionResult:
            enter("asr")
            try:
                message_id = int(prepared.path.stem.rsplit("-", 1)[1])
                if message_id == 1:
                    asr_started.set()
                    await overlap_observed.wait()
                await asyncio.sleep(0)
                return TranscriptionResult(text=f"Ответ {message_id}", provider="fake")
            finally:
                leave("asr")

    def message(message_id: int, user_id: int) -> SimpleNamespace:
        return SimpleNamespace(
            from_user=SimpleNamespace(id=user_id),
            chat=SimpleNamespace(id=100),
            message_id=message_id,
            date=datetime(2030, 1, 1, tzinfo=UTC),
            video_note=None,
            voice=None,
            audio=None,
            video=SimpleNamespace(
                file_id=f"file-{message_id}",
                file_unique_id=f"unique-{message_id}",
                file_name=f"feedback-{message_id}.mp4",
                mime_type="video/mp4",
                file_size=123,
            ),
            document=None,
        )

    monkeypatch.setattr("app.services.pipeline.download_media", fake_download_media)
    pipeline = ProcessingPipeline(
        settings,  # type: ignore[arg-type]
        repository,
        ControlledAudioConverter(),
        ControlledASR(),
    )

    async def run() -> list[object]:
        return await asyncio.wait_for(
            asyncio.gather(
                pipeline.process_media_message(message(1, 41), SimpleNamespace()),
                pipeline.process_media_message(message(2, 42), SimpleNamespace()),
                pipeline.process_media_message(message(3, 43), SimpleNamespace()),
            ),
            timeout=2,
        )

    results = asyncio.run(run())

    assert all(getattr(result, "ok") for result in results)
    assert overlap_snapshot == {"download": 1, "audio": 1, "asr": 1}
    assert stage_max == {"download": 3, "audio": 2, "asr": 1}