VLAarchtests/tests/test_rlbench_pickle_bootstrap.py

from __future__ import annotations

import sys

from sim_rlbench.dataset import _ensure_rlbench_pickle_modules


def test_rlbench_pickle_bootstrap_avoids_pyrep(monkeypatch) -> None:
    for key in list(sys.modules):
        if key == "pyrep" or key == "rlbench" or key.startswith("rlbench."):
            monkeypatch.delitem(sys.modules, key, raising=False)

    _ensure_rlbench_pickle_modules()

    assert "rlbench.demo" in sys.modules
    assert "rlbench.backend.observation" in sys.modules
    assert "pyrep" not in sys.modules