tests/test_scoring.py

"""Tests for dee/core/scoring.py — the shared ESM-2 scorer cache and
scoring concurrency guard used by both dee/server.py's /api/run pipeline
and dee/core/agent_tools.py's design_variant_library tool.
"""

import threading
import time

import pandas as pd
import pytest

from dee.core import scoring


@pytest.fixture(autouse=True)
def _semaphore_not_leaked():
    """Every test must return SCORING_SEMAPHORE to its initial (available)
    state — a leaked acquire would hang every later test using it."""
    yield
    acquired = scoring.SCORING_SEMAPHORE.acquire(timeout=0)
    assert acquired, "a test left SCORING_SEMAPHORE held"
    scoring.SCORING_SEMAPHORE.release()


@pytest.fixture(autouse=True)
def _clear_scorer_cache():
    scoring._SCORER_CACHE.clear()
    yield
    scoring._SCORER_CACHE.clear()


def test_effective_model_downgrades_gpu_models_without_gpu(monkeypatch):
    monkeypatch.setattr(scoring, "gpu_available", lambda: False)
    assert scoring.effective_model("medium") == "small"
    assert scoring.effective_model("large") == "small"
    assert scoring.effective_model("small") == "small"
    assert scoring.effective_model(None) == "small"


def test_effective_model_keeps_gpu_models_with_gpu(monkeypatch):
    monkeypatch.setattr(scoring, "gpu_available", lambda: True)
    assert scoring.effective_model("medium") == "medium"
    assert scoring.effective_model("large") == "large"


def test_get_scorer_caches_by_effective_config(monkeypatch):
    monkeypatch.setattr(scoring, "gpu_available", lambda: False)
    built = []

    class _FakeScorer:
        def __init__(self, config):
            built.append(config)

    monkeypatch.setattr("dee.models.scorer.ESM2Scorer", _FakeScorer)

    a = scoring.get_scorer("small")
    b = scoring.get_scorer("small")
    # 'medium' downgrades to 'small' on a no-GPU host — same cache key.
    c = scoring.get_scorer("medium")

    assert a is b is c
    assert len(built) == 1


def test_get_scorer_separate_cache_entries_per_model(monkeypatch):
    monkeypatch.setattr(scoring, "gpu_available", lambda: True)
    built = []

    class _FakeScorer:
        def __init__(self, config):
            built.append(config.model_name)

    monkeypatch.setattr("dee.models.scorer.ESM2Scorer", _FakeScorer)

    scoring.get_scorer("small")
    scoring.get_scorer("medium")
    assert built == ["small", "medium"]


def test_score_guarded_calls_scorer_and_returns_result():
    df = pd.DataFrame({"position": [0], "wt_aa": ["A"], "mut_aa": ["G"], "delta_ll": [1.0]})

    class _Scorer:
        def score_all_substitutions(self, protein):
            assert protein == "ACDE"
            return df

    result = scoring.score_guarded(_Scorer(), "ACDE")
    assert result is df


def test_score_guarded_raises_busy_when_semaphore_held():
    scoring.SCORING_SEMAPHORE.acquire()  # simulate another request scoring
    try:
        class _Scorer:
            def score_all_substitutions(self, protein):
                raise AssertionError("must not run while the semaphore is held")

        with pytest.raises(scoring.ScoringBusyError):
            scoring.score_guarded(_Scorer(), "ACDE", wait_timeout=0.05)
    finally:
        scoring.SCORING_SEMAPHORE.release()


def test_score_guarded_blocks_then_succeeds_once_released():
    scoring.SCORING_SEMAPHORE.acquire()

    def _release_soon():
        time.sleep(0.1)
        scoring.SCORING_SEMAPHORE.release()

    threading.Thread(target=_release_soon, daemon=True).start()

    class _Scorer:
        def score_all_substitutions(self, protein):
            return "scored"

    # No wait_timeout — blocks until the background thread releases, then
    # proceeds. If this hangs, the semaphore isn't being released/acquired
    # correctly.
    result = scoring.score_guarded(_Scorer(), "ACDE", wait_timeout=2.0)
    assert result == "scored"


def test_score_guarded_releases_semaphore_even_if_scorer_raises():
    class _Scorer:
        def score_all_substitutions(self, protein):
            raise RuntimeError("boom")

    with pytest.raises(RuntimeError):
        scoring.score_guarded(_Scorer(), "ACDE")
    # The autouse _semaphore_not_leaked fixture asserts this held true;
    # this test exists to name the specific behavior it's protecting.