tests/test_smolvla_runtime.py

from __future__ import annotations

import sys
import types
from collections import deque
from pathlib import Path

from dovla_cil.eval.smolvla_runtime import (
    ensure_lazy_policies_namespace,
    ensure_policy_utils_shim,
    ensure_train_config_type_shim,
)


def test_lazy_smolvla_namespace_points_to_installed_policy_directory(
    tmp_path: Path,
    monkeypatch,
) -> None:
    package_dir = tmp_path / "lerobot"
    policies_dir = package_dir / "policies"
    policies_dir.mkdir(parents=True)
    module = types.ModuleType("lerobot")
    module.__file__ = str(package_dir / "__init__.py")
    monkeypatch.delitem(sys.modules, "lerobot.policies", raising=False)

    namespace = ensure_lazy_policies_namespace(module)

    assert namespace.__path__ == [str(policies_dir)]
    assert namespace.__spec__ is not None
    assert namespace.__spec__.submodule_search_locations == [str(policies_dir)]
    assert sys.modules["lerobot.policies"] is namespace


def test_lazy_smolvla_namespace_is_idempotent(tmp_path: Path, monkeypatch) -> None:
    package_dir = tmp_path / "lerobot"
    (package_dir / "policies").mkdir(parents=True)
    module = types.ModuleType("lerobot")
    module.__file__ = str(package_dir / "__init__.py")
    monkeypatch.delitem(sys.modules, "lerobot.policies", raising=False)

    first = ensure_lazy_policies_namespace(module)
    second = ensure_lazy_policies_namespace(module)

    assert first is second


def test_train_config_shim_exposes_only_required_type(monkeypatch) -> None:
    configs = types.ModuleType("lerobot.configs")
    configs.__path__ = []
    monkeypatch.setitem(sys.modules, "lerobot.configs", configs)
    monkeypatch.delitem(sys.modules, "lerobot.configs.train", raising=False)

    module = ensure_train_config_type_shim()

    assert module.__dovla_type_only_shim__ is True
    assert module.TrainPipelineConfig.__module__ == "lerobot.configs.train"
    assert configs.train is module


def test_train_config_shim_preserves_real_import(monkeypatch) -> None:
    real_module = types.ModuleType("lerobot.configs.train")
    monkeypatch.setitem(sys.modules, "lerobot.configs.train", real_module)

    assert ensure_train_config_type_shim() is real_module


def test_policy_utils_shim_populates_and_advances_queues(monkeypatch) -> None:
    monkeypatch.delitem(sys.modules, "lerobot.policies.utils", raising=False)
    module = ensure_policy_utils_shim()
    queues = {"observation": deque(maxlen=2), "action": deque(maxlen=2)}

    result = module.populate_queues(
        queues,
        {"observation": "first", "action": "excluded"},
        exclude_keys=["action"],
    )
    module.populate_queues(queues, {"observation": "second"})

    assert result is queues
    assert list(queues["observation"]) == ["first", "second"]
    assert list(queues["action"]) == []


def test_policy_utils_shim_preserves_real_import(monkeypatch) -> None:
    real_module = types.ModuleType("lerobot.policies.utils")
    monkeypatch.setitem(sys.modules, "lerobot.policies.utils", real_module)

    assert ensure_policy_utils_shim() is real_module