tests/test_runtime.py

"""Step — compute config: cpu_threads resolution (auto / explicit cap / clamp)."""

from __future__ import annotations

import os

from common.runtime import _available_cores, configure_compute, onnx_session_options, resolve_cpu_threads


def test_auto_resolves_to_available_cores() -> None:
    assert resolve_cpu_threads(0) == _available_cores()
    assert resolve_cpu_threads(-3) == _available_cores()  # negative treated as auto, same as 0


def test_explicit_cap_below_available_is_kept() -> None:
    cap = max(1, _available_cores() - 1) if _available_cores() > 1 else 1
    assert resolve_cpu_threads(cap) == cap


def test_cap_above_available_is_clamped_down() -> None:
    huge = _available_cores() + 1000
    assert resolve_cpu_threads(huge) == _available_cores()


def test_configure_compute_sets_env_vars_to_resolved_value(monkeypatch) -> None:
    for var in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "NUMEXPR_NUM_THREADS",
                "POLARS_MAX_THREADS", "TOKENIZERS_PARALLELISM"):
        monkeypatch.delenv(var, raising=False)
    configure_compute(2)
    assert os.environ["OMP_NUM_THREADS"] == str(resolve_cpu_threads(2))
    assert os.environ["POLARS_MAX_THREADS"] == str(resolve_cpu_threads(2))
    assert os.environ["TOKENIZERS_PARALLELISM"] == "true"


def test_configure_compute_is_idempotent_setdefault(monkeypatch) -> None:
    monkeypatch.setenv("OMP_NUM_THREADS", "7")  # pre-existing value must survive (setdefault, not overwrite)
    configure_compute(2)
    assert os.environ["OMP_NUM_THREADS"] == "7"


def test_onnx_session_options_zero_is_auto() -> None:
    options = onnx_session_options(0)
    assert options.intra_op_num_threads == 0


def test_onnx_session_options_cap_is_passed_through() -> None:
    options = onnx_session_options(3)
    assert options.intra_op_num_threads == 3