tests/test_browsergym_lifecycle.py

import importlib.util
import sys
import threading
import types
from pathlib import Path

import pytest

from server.task_curriculum import TaskCurriculumScheduler


class FakeEnv:
    def __init__(self, name, close_error=None):
        self.name = name
        self.close_error = close_error
        self.closed = False
        self.reset_count = 0

    def close(self):
        if self.close_error is not None:
            raise self.close_error
        self.closed = True

    def reset(self, **_kwargs):
        self.reset_count += 1
        return {"url": f"https://example.test/{self.name}"}, {"goal": self.name}


class FakeState:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
        self.goal = kwargs.get("goal", "")
        self.current_url = kwargs.get("current_url", "")
        self.cum_reward = kwargs.get("cum_reward", 0.0)

    def model_copy(self, deep=False):
        return FakeState(**self.__dict__)


class FakeObservation:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)


class ReverseShuffle:
    def shuffle(self, values):
        values.reverse()


def load_browsergym_environment_module(monkeypatch):
    fake_interfaces = types.ModuleType("openenv.core.env_server.interfaces")
    fake_interfaces.Environment = object

    fake_models = types.ModuleType("browsergym_env.models")
    fake_models.BrowserGymAction = object
    fake_models.BrowserGymObservation = FakeObservation
    fake_models.BrowserGymState = FakeState

    fake_gymnasium = types.ModuleType("gymnasium")
    fake_gymnasium.make = lambda *_args, **_kwargs: FakeEnv("initial")

    monkeypatch.setitem(sys.modules, "openenv", types.ModuleType("openenv"))
    monkeypatch.setitem(sys.modules, "openenv.core", types.ModuleType("openenv.core"))
    monkeypatch.setitem(
        sys.modules,
        "openenv.core.env_server",
        types.ModuleType("openenv.core.env_server"),
    )
    monkeypatch.setitem(
        sys.modules, "openenv.core.env_server.interfaces", fake_interfaces
    )
    monkeypatch.setitem(sys.modules, "browsergym_env", types.ModuleType("browsergym_env"))
    monkeypatch.setitem(sys.modules, "browsergym_env.models", fake_models)
    monkeypatch.setitem(
        sys.modules, "browsergym_env.server", types.ModuleType("browsergym_env.server")
    )
    monkeypatch.setitem(
        sys.modules,
        "browsergym_env.server.task_curriculum",
        sys.modules["server.task_curriculum"],
    )
    monkeypatch.setitem(sys.modules, "gymnasium", fake_gymnasium)

    module_path = Path(__file__).resolve().parents[1] / "server" / "browsergym_environment.py"
    spec = importlib.util.spec_from_file_location(
        "browsergym_environment_under_test", module_path
    )
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def make_environment(module, task_name="click-test"):
    env = module.BrowserGymEnvironment.__new__(module.BrowserGymEnvironment)
    env._lock = threading.RLock()
    env.benchmark = "miniwob"
    env.task_name = task_name
    env.headless = True
    env.viewport_width = 1280
    env.viewport_height = 720
    env.timeout = 10000.0
    env.enable_task_curriculum = True
    env.task_scheduler = None
    env.gym_kwargs = {}
    env.env_id = env._build_env_id(task_name)
    env.gym_env = FakeEnv(task_name)
    env._state = FakeState(episode_id="old", step_count=3, benchmark="miniwob", task_name=task_name)
    env._last_obs = None
    env._last_info = None
    return env


def test_failed_gym_make_during_switch_leaves_old_env_intact(monkeypatch):
    module = load_browsergym_environment_module(monkeypatch)
    env = make_environment(module)
    old_env = env.gym_env
    old_state = env._state

    def fail_make(_env_id):
        raise ValueError("boom")

    env._make_gym_env = fail_make

    with pytest.raises(ValueError, match="boom"):
        env.reset(task_name="enter-text")

    assert env.gym_env is old_env
    assert old_env.closed is False
    assert env.task_name == "click-test"
    assert env.env_id == "browsergym/miniwob.click-test"
    assert env._state is old_state


def test_reset_switch_commits_after_successful_env_creation(monkeypatch):
    module = load_browsergym_environment_module(monkeypatch)
    env = make_environment(module)
    old_env = env.gym_env
    created_env = FakeEnv("enter-text")

    env._make_gym_env = lambda _env_id: created_env

    observation = env.reset(task_name="enter-text")

    assert env.gym_env is created_env
    assert old_env.closed is True
    assert env.task_name == "enter-text"
    assert env.env_id == "browsergym/miniwob.enter-text"
    assert env._state.task_name == "enter-text"
    assert observation.url == "https://example.test/enter-text"


def test_scheduler_selection_rolls_back_when_switch_creation_fails(monkeypatch):
    module = load_browsergym_environment_module(monkeypatch)
    env = make_environment(module)
    env.task_scheduler = TaskCurriculumScheduler(
        tasks=("enter-text",), rng=ReverseShuffle()
    )

    def fail_make(_env_id):
        raise ValueError("boom")

    env._make_gym_env = fail_make

    with pytest.raises(ValueError, match="boom"):
        env.reset()

    assert env.task_name == "click-test"
    assert env.env_id == "browsergym/miniwob.click-test"
    assert env.task_scheduler.next_task() == "enter-text"


def test_close_failure_during_switch_keeps_old_env_referenced(monkeypatch):
    module = load_browsergym_environment_module(monkeypatch)
    env = make_environment(module)
    close_error = RuntimeError("close failed")
    old_env = FakeEnv("click-test", close_error=close_error)
    new_env = FakeEnv("enter-text")
    env.gym_env = old_env
    env._make_gym_env = lambda _env_id: new_env

    with pytest.raises(RuntimeError, match="close failed"):
        env.reset(task_name="enter-text")

    assert env.gym_env is old_env
    assert new_env.closed is True
    assert env.task_name == "click-test"
    assert env.env_id == "browsergym/miniwob.click-test"


def test_state_returns_independent_snapshot(monkeypatch):
    module = load_browsergym_environment_module(monkeypatch)
    env = make_environment(module)

    state_snapshot = env.state
    state_snapshot.task_name = "mutated"
    state_snapshot.step_count = 999

    assert state_snapshot is not env._state
    assert env._state.task_name == "click-test"
    assert env._state.step_count == 3