Newer
Older
SmTIAS-LightSim / tests / runtime / test_overrides.py
"""runtime/overrides.py と,その runner / baker への組み込みの仕様ベーステスト.

仕様出典:
- docs/02_ENV/ENV_04_開発コマンド.md「環境変数」(SM_RENDER_VARIANT / SM_DRJIT_THREADS)
- 実装サマリー 仕様 7: 優先順位は --variant > SM_RENDER_VARIANT > YAML の render.variant
- 実装サマリー 仕様 8: SM_DRJIT_THREADS は runner / baker で dr.set_thread_count を適用する.
  未設定なら drjit を import しない.不正値(整数でない・1 未満)は ValueError

テスト方針:
- 環境変数は monkeypatch で設定・復元する
- 「drjit を import しない」は sys.modules["drjit"] = None(import すると ImportError)で検証する
- 実 drjit の set_thread_count も 1 本確かめ,元のスレッド数に戻す
"""

from __future__ import annotations

import logging
import sys
from pathlib import Path
from unittest.mock import MagicMock

import pytest

import sm_tias_lightsim.runtime.overrides as overrides
from sm_tias_lightsim.runtime.overrides import (
    THREADS_ENV_VAR,
    VARIANT_ENV_VAR,
    apply_drjit_thread_count,
    resolve_render_variant,
    resolve_thread_count,
)
from sm_tias_lightsim.scene.schema import SceneConfig


def test_env_var_names_match_spec():
    assert VARIANT_ENV_VAR == "SM_RENDER_VARIANT"
    assert THREADS_ENV_VAR == "SM_DRJIT_THREADS"


# ---------------------------------------------------------------------------
# resolve_render_variant(仕様 7)
# ---------------------------------------------------------------------------


class TestResolveRenderVariant:
    def test_yaml_is_used_when_nothing_else(self):
        assert resolve_render_variant("llvm_ad_rgb", None, {}) == "llvm_ad_rgb"

    def test_env_overrides_yaml(self):
        env = {VARIANT_ENV_VAR: "cuda_ad_rgb"}
        assert resolve_render_variant("llvm_ad_rgb", None, env) == "cuda_ad_rgb"

    def test_cli_overrides_env_and_yaml(self):
        env = {VARIANT_ENV_VAR: "cuda_ad_rgb"}
        assert resolve_render_variant("llvm_ad_rgb", "scalar_rgb", env) == "scalar_rgb"

    @pytest.mark.parametrize("blank", ["", "   "])
    def test_blank_cli_is_ignored(self, blank):
        env = {VARIANT_ENV_VAR: "cuda_ad_rgb"}
        assert resolve_render_variant("llvm_ad_rgb", blank, env) == "cuda_ad_rgb"

    @pytest.mark.parametrize("blank", ["", "   "])
    def test_blank_env_is_ignored(self, blank):
        assert (
            resolve_render_variant("llvm_ad_rgb", None, {VARIANT_ENV_VAR: blank}) == "llvm_ad_rgb"
        )

    def test_values_are_stripped(self):
        assert resolve_render_variant("llvm_ad_rgb", " cuda_ad_rgb ", {}) == "cuda_ad_rgb"
        env = {VARIANT_ENV_VAR: " scalar_rgb\n"}
        assert resolve_render_variant("llvm_ad_rgb", None, env) == "scalar_rgb"

    def test_defaults_to_process_environment(self, monkeypatch):
        monkeypatch.setenv(VARIANT_ENV_VAR, "cuda_ad_rgb")
        assert resolve_render_variant("llvm_ad_rgb") == "cuda_ad_rgb"
        monkeypatch.delenv(VARIANT_ENV_VAR)
        assert resolve_render_variant("llvm_ad_rgb") == "llvm_ad_rgb"


# ---------------------------------------------------------------------------
# resolve_thread_count(仕様 8)
# ---------------------------------------------------------------------------


class TestResolveThreadCount:
    @pytest.mark.parametrize("env", [{}, {THREADS_ENV_VAR: ""}, {THREADS_ENV_VAR: "   "}])
    def test_unset_returns_none(self, env):
        assert resolve_thread_count(env) is None

    @pytest.mark.parametrize(("raw", "expected"), [("1", 1), ("16", 16), (" 8 ", 8), ("064", 64)])
    def test_valid_values(self, raw, expected):
        assert resolve_thread_count({THREADS_ENV_VAR: raw}) == expected

    @pytest.mark.parametrize("raw", ["0", "-1", "abc", "1.5", "16 threads", "0x10"])
    def test_invalid_values_raise_value_error(self, raw):
        """整数でない・1 未満は ValueError(黙って全コアで走らせない)."""
        with pytest.raises(ValueError, match=THREADS_ENV_VAR):
            resolve_thread_count({THREADS_ENV_VAR: raw})

    def test_defaults_to_process_environment(self, monkeypatch):
        monkeypatch.setenv(THREADS_ENV_VAR, "12")
        assert resolve_thread_count() == 12


# ---------------------------------------------------------------------------
# apply_drjit_thread_count(仕様 8)
# ---------------------------------------------------------------------------


@pytest.fixture
def reset_applied(monkeypatch):
    monkeypatch.setattr(overrides, "_applied_thread_count", None)


class TestApplyDrjitThreadCount:
    def test_unset_does_not_import_drjit(self, monkeypatch, reset_applied):
        """未設定なら何もせず,drjit を import しない."""
        monkeypatch.setitem(sys.modules, "drjit", None)  # import されると ImportError になる
        assert apply_drjit_thread_count({}) is None

    def test_sets_thread_count(self, monkeypatch, reset_applied):
        fake_dr = MagicMock()
        monkeypatch.setitem(sys.modules, "drjit", fake_dr)
        assert apply_drjit_thread_count({THREADS_ENV_VAR: "16"}) == 16
        fake_dr.set_thread_count.assert_called_once_with(16)

    def test_invalid_value_raises_before_touching_drjit(self, monkeypatch, reset_applied):
        fake_dr = MagicMock()
        monkeypatch.setitem(sys.modules, "drjit", fake_dr)
        with pytest.raises(ValueError):
            apply_drjit_thread_count({THREADS_ENV_VAR: "0"})
        fake_dr.set_thread_count.assert_not_called()

    def test_reads_process_environment(self, monkeypatch, reset_applied):
        fake_dr = MagicMock()
        monkeypatch.setitem(sys.modules, "drjit", fake_dr)
        monkeypatch.setenv(THREADS_ENV_VAR, "3")
        assert apply_drjit_thread_count() == 3
        fake_dr.set_thread_count.assert_called_once_with(3)

    def test_info_log_once_for_repeated_calls(self, monkeypatch, reset_applied, caplog):
        """runner と baker の両方から呼ばれても同じ値なら INFO は 1 回(冪等)."""
        fake_dr = MagicMock()
        monkeypatch.setitem(sys.modules, "drjit", fake_dr)
        with caplog.at_level(logging.INFO, logger=overrides.__name__):
            apply_drjit_thread_count({THREADS_ENV_VAR: "4"})
            apply_drjit_thread_count({THREADS_ENV_VAR: "4"})
        infos = [r for r in caplog.records if r.levelno == logging.INFO]
        assert len(infos) == 1
        assert fake_dr.set_thread_count.call_count == 2

    def test_real_drjit_thread_count_is_applied(self, reset_applied):
        """実物の Dr.Jit でスレッド数が変わること(元に戻す)."""
        import drjit as dr

        original = dr.thread_count()
        target = 2 if original != 2 else 3
        try:
            assert apply_drjit_thread_count({THREADS_ENV_VAR: str(target)}) == target
            assert dr.thread_count() == target
        finally:
            dr.set_thread_count(original)


# ---------------------------------------------------------------------------
# runner / baker への組み込み(仕様 8)
# ---------------------------------------------------------------------------


class _StopError(Exception):
    """set_thread_count が呼ばれた時点で処理を打ち切るための例外."""


def _scene_config(variant: str = "scalar_rgb") -> SceneConfig:
    return SceneConfig.model_validate(
        {
            "name": "threads",
            "geometry": {
                "meshes": [{"id": "m", "path": "dummy.stl", "unit": "mm", "material": "white"}]
            },
            "materials": {"white": {"type": "diffuse", "reflectance": 0.9}},
            "lights": [
                {
                    "id": "led",
                    "type": "point",
                    "position": [0.0, 0.0, 100.0],
                    "intensity": 1.0,
                    "profile": "lambertian",
                }
            ],
            "sensor": {
                "type": "perspective",
                "position": [0.0, 0.0, 200.0],
                "target": [0.0, 0.0, 0.0],
                "up": [0.0, 1.0, 0.0],
                "fov_deg": 60.0,
                "film": {"width": 4, "height": 4},
            },
            "render": {"variant": variant, "integrator": "path", "max_depth": 2, "spp": 1},
        }
    )


class TestRunnerAppliesThreadCount:
    @pytest.fixture(autouse=True)
    def _isolate_variant(self, monkeypatch, reset_applied):
        import sm_tias_lightsim.render.runner as runner

        monkeypatch.setattr(runner, "_current_variant", None)

    def test_render_scene_applies_env_thread_count(self, monkeypatch):
        from sm_tias_lightsim.render.runner import render_scene

        fake_dr = MagicMock()
        fake_dr.set_thread_count.side_effect = _StopError
        monkeypatch.setitem(sys.modules, "drjit", fake_dr)
        monkeypatch.setenv(THREADS_ENV_VAR, "5")
        with pytest.raises(_StopError):
            render_scene(_scene_config(), project_root=Path("."))
        fake_dr.set_thread_count.assert_called_once_with(5)

    def test_render_scene_rejects_invalid_env(self, monkeypatch):
        from sm_tias_lightsim.render.runner import render_scene

        monkeypatch.setenv(THREADS_ENV_VAR, "abc")
        with pytest.raises(ValueError, match=THREADS_ENV_VAR):
            render_scene(_scene_config(), project_root=Path("."))


class TestBakerAppliesThreadCount:
    @pytest.mark.parametrize("mode", ["direct", "path"])
    def test_bake_irradiance_applies_env_thread_count(self, monkeypatch, reset_applied, mode):
        from sm_tias_lightsim.visualize.baker import bake_irradiance

        fake_dr = MagicMock()
        fake_dr.set_thread_count.side_effect = _StopError
        monkeypatch.setitem(sys.modules, "drjit", fake_dr)
        monkeypatch.setenv(THREADS_ENV_VAR, "6")
        with pytest.raises(_StopError):
            bake_irradiance(_scene_config(), Path("."), mode=mode)
        fake_dr.set_thread_count.assert_called_once_with(6)

    def test_bake_irradiance_rejects_invalid_env(self, monkeypatch, reset_applied):
        from sm_tias_lightsim.visualize.baker import bake_irradiance

        monkeypatch.setenv(THREADS_ENV_VAR, "-2")
        with pytest.raises(ValueError, match=THREADS_ENV_VAR):
            bake_irradiance(_scene_config(), Path("."), mode="direct")