talkscore-asr/tests/test_concurrency.py

"""Тесты параллельной обработки.

Замеры показали: обе стадии упираются в 4 потока, а на 16 работают вчетверо
медленнее. Значит ядра нужно занимать не шириной одной задачи, а несколькими
задачами сразу - и тогда очередь обязана быть устойчивой к гонкам.
"""
import threading

import pytest

from app.store import JobStatus, JobStore


@pytest.fixture
def store(tmp_path):
    return JobStore(tmp_path / "jobs.db")


class TestClaimIsAtomic:
    def test_claim_marks_running(self, store):
        job_id = store.create(filename="a.wav", duration_sec=1.0)
        assert store.claim_next() == job_id
        assert store.get(job_id)["status"] == JobStatus.RUNNING

    def test_second_claim_gets_nothing(self, store):
        store.create(filename="a.wav", duration_sec=1.0)
        store.claim_next()
        assert store.claim_next() is None

    def test_each_job_claimed_once_under_load(self, store):
        """Главное требование: два воркера не должны взять одну задачу."""
        ids = {store.create(filename=f"{i}.wav", duration_sec=1.0) for i in range(50)}
        claimed: list[str] = []
        lock = threading.Lock()

        def worker():
            while True:
                job_id = store.claim_next()
                if job_id is None:
                    return
                with lock:
                    claimed.append(job_id)

        threads = [threading.Thread(target=worker) for _ in range(8)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        assert len(claimed) == len(set(claimed)) == 50
        assert set(claimed) == ids

    def test_claims_oldest_first(self, store):
        first = store.create(filename="1.wav", duration_sec=1.0)
        store.create(filename="2.wav", duration_sec=1.0)
        assert store.claim_next() == first


class TestWorkerSettings:
    def test_default_threads_is_four_not_all_cores(self, tmp_path):
        """Широкие потоки замедляют обе стадии, поэтому по умолчанию их немного."""
        from app.config import load_settings

        config = tmp_path / "config.toml"
        config.write_text('[processing]\nthreads=0\n', encoding="utf-8")
        assert load_settings(config).effective_threads() == 4

    def test_explicit_threads_respected(self, tmp_path):
        from app.config import load_settings

        config = tmp_path / "config.toml"
        config.write_text('[processing]\nthreads=6\n', encoding="utf-8")
        assert load_settings(config).effective_threads() == 6

    def test_workers_derived_from_cores(self, tmp_path, monkeypatch):
        from app.config import load_settings

        monkeypatch.setattr("os.cpu_count", lambda: 32)
        config = tmp_path / "config.toml"
        config.write_text('[processing]\nthreads=4\nworkers=0\n', encoding="utf-8")
        # 32 логических ядра при 4 потоках на задачу - но не больше разумного предела
        assert 2 <= load_settings(config).effective_workers() <= 4

    def test_workers_never_below_one(self, tmp_path, monkeypatch):
        from app.config import load_settings

        monkeypatch.setattr("os.cpu_count", lambda: 1)
        config = tmp_path / "config.toml"
        config.write_text('[processing]\nthreads=4\nworkers=0\n', encoding="utf-8")
        assert load_settings(config).effective_workers() == 1

    def test_explicit_workers_respected(self, tmp_path):
        from app.config import load_settings

        config = tmp_path / "config.toml"
        config.write_text('[processing]\nworkers=3\n', encoding="utf-8")
        assert load_settings(config).effective_workers() == 3


class TestLazyWarmup:
    """Воркеры создаются заранее, а модели грузят при первой своей задаче."""

    def test_transcribe_warms_up_when_needed(self, tmp_path, monkeypatch):
        import sys
        import types

        for name in ("sherpa_onnx", "onnx_asr", "onnxruntime"):
            monkeypatch.setitem(sys.modules, name, types.ModuleType(name))
        from app.pipeline import Pipeline

        p = Pipeline(models_dir=tmp_path, threads=1,
                     replacements_path=tmp_path / "r.txt", base_dir=tmp_path)
        called = {"warmup": 0}
        monkeypatch.setattr(p, "warmup", lambda: called.__setitem__("warmup", 1))
        # transcribe упадёт дальше на чтении файла, но warmup обязан быть вызван
        with pytest.raises(Exception):
            p.transcribe(tmp_path / "нет.wav")
        assert called["warmup"] == 1

    def test_no_runtime_error_about_warmup(self, tmp_path, monkeypatch):
        """Прежде здесь падало «модели не загружены, вызовите warmup()»."""
        import sys
        import types

        for name in ("sherpa_onnx", "onnx_asr", "onnxruntime"):
            monkeypatch.setitem(sys.modules, name, types.ModuleType(name))
        from app.pipeline import Pipeline

        p = Pipeline(models_dir=tmp_path, threads=1,
                     replacements_path=tmp_path / "r.txt", base_dir=tmp_path)
        monkeypatch.setattr(p, "warmup", lambda: None)
        try:
            p.transcribe(tmp_path / "нет.wav")
        except RuntimeError as exc:
            assert "warmup" not in str(exc)
        except Exception:
            pass


class TestBenchmarkEndpoint:
    def test_benchmark_requires_token(self, tmp_path, monkeypatch):
        import sys
        import types

        from fastapi.testclient import TestClient

        config = tmp_path / "config.toml"
        config.write_text('[security]\ntoken="t"\nallow_ips=""\n', encoding="utf-8")
        monkeypatch.setenv("TALKSCORE_ASR_HOME", str(tmp_path))
        for name in ("sherpa_onnx", "onnx_asr", "onnxruntime"):
            monkeypatch.setitem(sys.modules, name, types.ModuleType(name))
        for mod in [m for m in list(sys.modules) if m.startswith("app.")]:
            del sys.modules[mod]
        import app.config as cfg
        monkeypatch.setattr(cfg, "BASE_DIR", tmp_path)
        import app.main as main
        main._worker_stop.set()

        with TestClient(main.app) as c:
            assert c.post("/v1/benchmark", files={"file": ("a.wav", b"x")}).status_code == 401