From 10d0547699f3cea73647b3aabecff7182a6b0c35 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 14 Oct 2026 15:37:04 +0000 Subject: [PATCH] test: feature - hermetic unit tests for experiments loop, exp_utils, launch_exp and ui_assistant Co-Authored-By: Daniella Grimberg --- tests/experiments/test_exp_utils_unit.py | 190 +++++ tests/experiments/test_launch_exp_unit.py | 324 ++++++++ tests/experiments/test_loop_unit.py | 935 ++++++++++++++++++++++ tests/test_ui_assistant_unit.py | 121 +++ 4 files changed, 1570 insertions(+) create mode 100644 tests/experiments/test_exp_utils_unit.py create mode 100644 tests/experiments/test_launch_exp_unit.py create mode 100644 tests/experiments/test_loop_unit.py create mode 100644 tests/test_ui_assistant_unit.py diff --git a/tests/experiments/test_exp_utils_unit.py b/tests/experiments/test_exp_utils_unit.py new file mode 100644 index 000000000..e5c66b042 --- /dev/null +++ b/tests/experiments/test_exp_utils_unit.py @@ -0,0 +1,190 @@ +"""Hermetic unit tests for agentlab.experiments.exp_utils.""" + +import importlib.util +import signal +import sys +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + +import pytest + +from agentlab.experiments import exp_utils +from agentlab.experiments.exp_utils import ( + MockedExpArgs, + _episode_timeout, + add_dependencies, + run_exp, + timeout_manager, +) +from agentlab.experiments.loop import EnvArgs, ExpArgs + +# --------------------------------------------------------------------------- RESULTS_DIR + + +def _load_fresh_copy(name): + """Execute exp_utils' source in an isolated module object (doesn't touch sys.modules).""" + spec = importlib.util.spec_from_file_location(name, exp_utils.__file__) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize( + "agentlab_root, copilot_root, expected", + [ + ("agentlab_root", None, "agentlab_root"), + ("agentlab_root", "copilot_root", "agentlab_root"), + (None, "copilot_root", "copilot_root"), + (None, None, "home/agentlab_results"), + ], +) +def test_results_dir_resolution(tmp_path, monkeypatch, agentlab_root, copilot_root, expected): + monkeypatch.setenv("HOME", str(tmp_path / "home")) + for var, value in [ + ("AGENTLAB_EXP_ROOT", agentlab_root), + ("UI_COPILOT_RESULTS_DIR", copilot_root), + ]: + if value is None: + monkeypatch.delenv(var, raising=False) + else: + monkeypatch.setenv(var, str(tmp_path / value)) + + module = _load_fresh_copy("_exp_utils_results_dir_copy") + assert isinstance(module.RESULTS_DIR, Path) + assert module.RESULTS_DIR == tmp_path / expected + assert module.RESULTS_DIR.is_dir() + + +# --------------------------------------------------------------------------- run_exp / timeouts + + +def test_run_exp_delegates_to_run(): + exp = mock.Mock() + exp.run.return_value = "result" + assert run_exp(exp, "dep1", "dep2", avg_step_timeout=1) == "result" + exp.run.assert_called_once_with() + + +@pytest.mark.parametrize( + "max_steps, avg, episode_timeout, expected", + [ + (None, 60, None, 10 * 60 * 60), + (5, 10, None, 50), + (5, 10, 20, 20), + (5, 10, 500, 50), + (None, 60, 100, 100), + ], +) +def test_episode_timeout(max_steps, avg, episode_timeout, expected): + exp = SimpleNamespace(env_args=SimpleNamespace(max_steps=max_steps)) + if episode_timeout is not None: + exp.episode_timeout = episode_timeout + assert _episode_timeout(exp, avg_step_timeout=avg) == expected + + +def test_episode_timeout_env_args_without_max_steps(): + assert _episode_timeout(SimpleNamespace(env_args=object())) == 10 * 60 * 60 + + +def test_timeout_manager_none_is_noop(monkeypatch): + alarm = mock.Mock() + monkeypatch.setattr(exp_utils.signal, "alarm", alarm) + with timeout_manager(None): + pass + alarm.assert_not_called() + + +def test_timeout_manager_windows_is_noop(monkeypatch): + alarm = mock.Mock() + monkeypatch.setattr(exp_utils.signal, "alarm", alarm) + monkeypatch.setattr(exp_utils.sys, "platform", "win32") + with timeout_manager(5): + pass + alarm.assert_not_called() + + +@pytest.mark.parametrize("seconds, expected_alarm", [(5, 5), (0.2, 1), (3.7, 3)]) +def test_timeout_manager_sets_and_restores_alarm(monkeypatch, seconds, expected_alarm): + alarm = mock.Mock() + monkeypatch.setattr(exp_utils.signal, "alarm", alarm) + previous = signal.getsignal(signal.SIGALRM) + with timeout_manager(seconds): + handler = signal.getsignal(signal.SIGALRM) + assert handler is not previous and callable(handler) + assert alarm.call_args_list == [mock.call(expected_alarm), mock.call(0)] + assert signal.getsignal(signal.SIGALRM) == previous + + +def test_timeout_manager_handler_raises_timeout(monkeypatch): + alarm = mock.Mock() + monkeypatch.setattr(exp_utils.signal, "alarm", alarm) + previous = signal.getsignal(signal.SIGALRM) + with pytest.raises(TimeoutError, match="timed out after 4 seconds"): + with timeout_manager(4): + signal.getsignal(signal.SIGALRM)(signal.SIGALRM, None) # simulate the alarm firing + assert alarm.call_args_list[-1] == mock.call(0) + assert signal.getsignal(signal.SIGALRM) == previous + + +# --------------------------------------------------------------------------- add_dependencies + + +def _exp(task_name, exp_id=None): + return ExpArgs(agent_args=None, env_args=EnvArgs(task_name=task_name), exp_id=exp_id) + + +@pytest.mark.parametrize("deps", [None, {}, {"a": [], "b": []}]) +def test_add_dependencies_noop(deps): + exps = [_exp("a"), _exp("b")] + assert add_dependencies(exps, deps) is exps + assert all(e.exp_id is None and e.depends_on == () for e in exps) + + +def test_add_dependencies_maps_task_names_to_ids(): + exps = [_exp("a", "id-a"), _exp("b"), _exp("c", "id-c")] + out = add_dependencies(exps, {"a": [], "b": ["a"], "c": ["a", "b"]}) + assert out is exps + id_b = exps[1].exp_id + assert id_b is not None + assert exps[0].depends_on == () + assert exps[1].depends_on == ("id-a",) + assert exps[2].depends_on == ("id-a", id_b) + + +def test_add_dependencies_rejects_duplicate_tasks(): + with pytest.raises(ValueError, match="not unique"): + add_dependencies([_exp("a"), _exp("a")], {"a": ["b"]}) + + +def test_add_dependencies_rejects_missing_task(): + with pytest.raises(ValueError, match="Task b is missing"): + add_dependencies([_exp("a"), _exp("b")], {"a": ["b"]}) + + +# --------------------------------------------------------------------------- MockedExpArgs & deprecated + + +def test_mocked_exp_args(monkeypatch): + sleep = mock.Mock() + monkeypatch.setattr(exp_utils, "sleep", sleep) + monkeypatch.setattr(exp_utils, "time", mock.Mock(side_effect=[100.0, 103.0])) + m = MockedExpArgs("t1", depends_on=["t0"]) + assert (m.exp_name, m.depends_on, m.env_args) == ("exp_t1", ["t0"], None) + assert m.run() is m + sleep.assert_called_once_with(3) + assert (m.start_time, m.end_time) == (100.0, 103.0) + assert MockedExpArgs("t2").depends_on == [] + + +@pytest.mark.parametrize( + "fn, args", + [ + (exp_utils.make_seeds, (3,)), + (exp_utils.order, ([],)), + (exp_utils.hide_some_exp, ("dir", lambda e: True, True)), + ], +) +def test_deprecated_helpers_raise(fn, args): + with pytest.raises(DeprecationWarning, match="will be removed"): + fn(*args) diff --git a/tests/experiments/test_launch_exp_unit.py b/tests/experiments/test_launch_exp_unit.py new file mode 100644 index 000000000..fa8302d0b --- /dev/null +++ b/tests/experiments/test_launch_exp_unit.py @@ -0,0 +1,324 @@ +"""Hermetic unit tests for agentlab.experiments.launch_exp (no browser, no real Ray).""" + +import json +import os +import sys +import types +from dataclasses import dataclass, field +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + +import gymnasium as gym +import pytest +from browsergym.experiments.agent import Agent + +from agentlab.experiments import launch_exp, loop +from agentlab.experiments.launch_exp import ( + _hide_completed, + _split_path, + _split_sequential_exp, + find_incomplete, + import_object, + non_dummy_count, + noop, + run_experiments, +) +from agentlab.experiments.loop import AbstractAgentArgs, EnvArgs, ExpArgs, ExpResult + + +class _ActionSet: + def to_python_code(self, action): + return action + + +class _Agent(Agent): + action_set = _ActionSet() + + def obs_preprocessor(self, obs): + return obs + + def get_action(self, obs): + return "noop()", {} + + +@dataclass +class TrackingAgentArgs(AbstractAgentArgs): + events: list = field(default_factory=list, repr=False) + + def prepare(self): + self.events.append("prepare") + + def close(self): + self.events.append("close") + + def make_agent(self): + return _Agent() + + +class _OneStepEnv(gym.Env): + def reset(self, seed=None, options=None): + return {"goal": "g"}, {} + + def step(self, action): + info = {"action_exec_start": 0.0, "action_exec_stop": 1.0, "action_exec_timeout": 0.0} + return {"goal": "g"}, 1.0, True, False, info + + +@pytest.fixture(autouse=True) +def isolate(monkeypatch): + monkeypatch.setattr(loop, "EXP_RESULT_CACHE", {}) + monkeypatch.setattr(loop, "_is_debugging", lambda: False) + # ExpArgs.run() reconfigures the root logger; keep it from leaking between tests. + monkeypatch.setattr(ExpArgs, "_set_logger", lambda self: None) + monkeypatch.setattr(ExpArgs, "_unset_logger", lambda self: None) + + +def _exps(n, **kwargs): + return [ + ExpArgs( + agent_args=TrackingAgentArgs(), + env_args=EnvArgs(task_name=f"task{i}", task_seed=i), + **kwargs, + ) + for i in range(n) + ] + + +@pytest.fixture +def recorded_runs(monkeypatch): + runs = [] + + def fake_run_exp(exp_args, *deps, avg_step_timeout=60): + runs.append((exp_args.env_args.task_name, avg_step_timeout)) + + monkeypatch.setattr(launch_exp, "run_exp", fake_run_exp) + return runs + + +# --------------------------------------------------------------------------- run_experiments + + +def test_run_experiments_empty_list(tmp_path): + study_dir = tmp_path / "study" + assert run_experiments(1, [], study_dir, parallel_backend="sequential") is None + assert not study_dir.exists() + + +def test_run_experiments_sequential_end_to_end(tmp_path, monkeypatch): + monkeypatch.setattr(EnvArgs, "make_env", lambda self, **kw: _OneStepEnv()) + exps = _exps(2) + study_dir = tmp_path / "nested" / "study" + run_experiments(1, exps, study_dir, parallel_backend="sequential") + + assert study_dir.is_dir() + for exp in exps: + assert exp.agent_args.events == ["prepare", "close"] + assert Path(exp.exp_dir).parent == study_dir + summary = json.loads((Path(exp.exp_dir) / "summary_info.json").read_text()) + assert summary["cum_reward"] == 1.0 and summary["err_msg"] is None + assert ExpResult(exp.exp_dir).status == "done" + + +def test_run_experiments_sequential_passes_timeout(tmp_path, recorded_runs): + run_experiments(1, _exps(3), tmp_path, parallel_backend="sequential", avg_step_timeout=7) + assert recorded_runs == [("task0", 7), ("task1", 7), ("task2", 7)] + + +class _FakeParallel: + instances = [] + + def __init__(self, n_jobs, prefer): + self.n_jobs, self.prefer = n_jobs, prefer + _FakeParallel.instances.append(self) + + def __call__(self, jobs): + return [fn(*args, **kwargs) for fn, args, kwargs in jobs] + + +def test_run_experiments_joblib_runs_sequential_first(tmp_path, monkeypatch, recorded_runs): + import joblib + + monkeypatch.setattr(_FakeParallel, "instances", []) + monkeypatch.setattr(joblib, "Parallel", _FakeParallel) + exps = _exps(3) + exps[2].sequential = True + + run_experiments(4, exps, tmp_path, parallel_backend="joblib", avg_step_timeout=5) + + assert recorded_runs == [("task2", 5), ("task0", 5), ("task1", 5)] + (par,) = _FakeParallel.instances + assert (par.n_jobs, par.prefer) == (4, "processes") + assert [e.agent_args.events[0] for e in exps] == ["prepare"] * 3 + + +@pytest.mark.xfail( + strict=True, + reason="BUG: joblib branch reassigns exp_args_list, so sequential agents never get close()", +) +def test_run_experiments_joblib_closes_all_agents(tmp_path, monkeypatch, recorded_runs): + import joblib + + monkeypatch.setattr(joblib, "Parallel", _FakeParallel) + exps = _exps(2) + exps[0].sequential = True + run_experiments(2, exps, tmp_path, parallel_backend="joblib") + assert [e.agent_args.events for e in exps] == [["prepare", "close"]] * 2 + + +@pytest.fixture +def fake_ray_module(monkeypatch): + mod = types.ModuleType("agentlab.experiments.graph_execution_ray") + mod.ray = mock.Mock() + mod.execute_task_graph = mock.Mock() + monkeypatch.setitem(sys.modules, "agentlab.experiments.graph_execution_ray", mod) + return mod + + +def test_run_experiments_ray_backend_with_fake_ray(tmp_path, fake_ray_module): + exps = _exps(2) + run_experiments(3, exps, tmp_path, parallel_backend="ray", avg_step_timeout=9) + fake_ray_module.ray.init.assert_called_once_with(num_cpus=3) + fake_ray_module.execute_task_graph.assert_called_once_with(exps, avg_step_timeout=9) + fake_ray_module.ray.shutdown.assert_called_once() + + +def test_run_experiments_ray_shutdown_on_error(tmp_path, fake_ray_module): + fake_ray_module.execute_task_graph.side_effect = RuntimeError("worker died") + exps = _exps(1) + with pytest.raises(RuntimeError, match="worker died"): + run_experiments(1, exps, tmp_path, parallel_backend="ray") + fake_ray_module.ray.shutdown.assert_called_once() + assert exps[0].agent_args.events == ["prepare", "close"] + + +def test_run_experiments_unknown_backend_still_closes(tmp_path, recorded_runs): + exps = _exps(2) + with pytest.raises(ValueError, match="Unknown parallel_backend: dask"): + run_experiments(1, exps, tmp_path, parallel_backend="dask") + assert recorded_runs == [] + assert all(e.agent_args.events == ["prepare", "close"] for e in exps) + # experiments were prepared on disk before the backend was chosen + assert all((Path(e.exp_dir) / "exp_args.pkl").exists() for e in exps) + + +# --------------------------------------------------------------------------- find_incomplete + + +def _make_study(study_dir, statuses): + """Create prepared experiments with the given status ('done'/'error'/'incomplete').""" + exps = [] + for i, status in enumerate(statuses): + exp = ExpArgs( + agent_args=TrackingAgentArgs(), + env_args=EnvArgs(task_name=f"task{i}", task_seed=0), + order=len(statuses) - 1 - i, + ) + exp.prepare(study_dir) + if status != "incomplete": + summary = {"terminated": True, "err_msg": "boom" if status == "error" else None} + (Path(exp.exp_dir) / "summary_info.json").write_text(json.dumps(summary)) + exps.append(exp) + return exps + + +def test_find_incomplete_missing_dir(tmp_path): + with pytest.raises(ValueError, match="does not exist"): + find_incomplete(tmp_path / "nope") + + +@pytest.mark.parametrize("include_errors, expected_count", [(True, 2), (False, 1)]) +def test_find_incomplete(tmp_path, include_errors, expected_count): + _make_study(tmp_path, ["done", "error", "incomplete"]) + exps = find_incomplete(tmp_path, include_errors=include_errors) + + assert non_dummy_count(exps) == expected_count + # sorted by `order`, which was assigned in reverse creation order + assert [e.env_args.task_name for e in exps] == ["task2", "task1", "task0"] + assert [e.status for e in exps] == ["incomplete", "error", "done"] + done = exps[2] + assert done.is_dummy and done.run is noop and done.prepare is noop + assert exps[0].is_dummy is False + assert exps[1].is_dummy is (not include_errors) + + +def test_find_incomplete_all_done(tmp_path): + _make_study(tmp_path, ["done", "done"]) + exps = find_incomplete(tmp_path) + assert len(exps) == 2 and non_dummy_count(exps) == 0 + + +def test_find_incomplete_orders_none_as_zero(tmp_path): + exps = _make_study(tmp_path, ["incomplete", "incomplete"]) + for exp in exps: # rewrite pickles with order=None for the first one + exp.order = None if exp.env_args.task_name == "task0" else 1 + with open(Path(exp.exp_dir) / "exp_args.pkl", "wb") as f: + import pickle + + pickle.dump(exp, f) + found = find_incomplete(tmp_path) + assert [e.env_args.task_name for e in found] == ["task0", "task1"] + + +def test_hide_completed_with_stub_result(): + exp_args = SimpleNamespace(run="run", prepare="prepare") + res = SimpleNamespace(status="error", exp_args=exp_args) + out = _hide_completed(res, include_errors=False) + assert out is exp_args and out.is_dummy and out.run is noop and out.status == "error" + + +# --------------------------------------------------------------------------- small helpers + + +def test_non_dummy_count(): + exps = [SimpleNamespace(is_dummy=d) for d in (True, False, False)] + assert non_dummy_count(exps) == 2 + assert non_dummy_count([]) == 0 + + +def test_noop(): + assert noop() is None + assert noop(1, 2, a=3) is None + + +def test_split_sequential_exp(): + a, b, c = SimpleNamespace(sequential=True), SimpleNamespace(), SimpleNamespace(sequential=False) + assert _split_sequential_exp([a, b, c]) == ([a], [b, c]) + assert _split_sequential_exp([]) == ([], []) + + +@pytest.mark.parametrize( + "path, expected", + [ + ("pkg.mod.obj", ("pkg.mod", "obj")), + ("pkg/mod/obj", ("pkg.mod", "obj")), + ("mod.obj", ("mod", "obj")), + ], +) +def test_split_path(path, expected): + assert _split_path(path) == expected + + +def test_split_path_requires_a_dot(): + with pytest.raises(ValueError): + _split_path("noseparator") + + +@pytest.mark.parametrize( + "path, expected", + [ + ("agentlab.experiments.launch_exp.noop", noop), + ("agentlab/experiments/launch_exp/noop", noop), + ("os.path.join", os.path.join), + ], +) +def test_import_object(path, expected): + assert import_object(path) is expected + + +@pytest.mark.parametrize( + "path", ["agentlab.does_not_exist_xyz.obj", "agentlab.experiments.launch_exp.nope"] +) +def test_import_object_errors(path): + with pytest.raises(ImportError, match=f"Error importing {path}"): + import_object(path) diff --git a/tests/experiments/test_loop_unit.py b/tests/experiments/test_loop_unit.py new file mode 100644 index 000000000..d63150974 --- /dev/null +++ b/tests/experiments/test_loop_unit.py @@ -0,0 +1,935 @@ +"""Hermetic unit tests for agentlab.experiments.loop (fake env + fake agent, no browser).""" + +import gzip +import json +import logging +import pickle +import sys +import types +from dataclasses import dataclass, field +from pathlib import Path +from unittest import mock + +import gymnasium as gym +import numpy as np +import pytest +from browsergym.core.chat import Chat +from browsergym.experiments.agent import Agent +from PIL import Image + +from agentlab.experiments import loop +from agentlab.experiments.loop import ( + AbstractAgentArgs, + DataclassJSONEncoder, + EnvArgs, + ExpArgs, + ExpResult, + StepInfo, + StepTimestamps, +) + +# --------------------------------------------------------------------------- fakes + + +class FakeActionSet: + def to_python_code(self, action): + return f"# {action}" + + +class FakeAgent(Agent): + action_set = FakeActionSet() + + def __init__(self, actions, raise_at=None, exc=ValueError, think="thinking"): + self.actions = list(actions) + self.raise_at = raise_at + self.exc = exc + self.think = think + self.calls = 0 + self.task_name = None + self.final_tape = "the-tape" + + def set_task_name(self, task_name): + self.task_name = task_name + + def obs_preprocessor(self, obs): + obs = dict(obs) + obs["preprocessed"] = True + return obs + + def get_action(self, obs): + idx = self.calls + self.calls += 1 + if self.raise_at is not None and idx == self.raise_at: + raise self.exc("agent exploded") + action = self.actions[idx] if idx < len(self.actions) else None + info = {"stats": {"llm_cost": 0.5}} + if self.think is not None: + info["think"] = self.think + return action, info + + +@dataclass +class FakeAgentArgs(AbstractAgentArgs): + actions: tuple = ("click('a')", "click('b')") + raise_at: int = None + agents: list = field(default_factory=list, repr=False) + + def make_agent(self): + agent = FakeAgent(self.actions, raise_at=self.raise_at) + self.agents.append(agent) + return agent + + +class FakeEnv(gym.Env): + """Episode terminates after `n_steps` env.step() calls.""" + + def __init__(self, n_steps=2, with_chat=False, goal_object=True): + self.n_steps = n_steps + self.t = 0 + self.closed = False + self.actions = [] + self.reset_seed = None + self.goal_object = goal_object + if with_chat: + self.chat = mock.MagicMock(spec=Chat) + + def _obs(self): + obs = { + "goal": "reach the end", + "axtree_txt": "button 'ok'", + "screenshot": np.full((4, 4, 3), self.t * 10, dtype=np.uint8), + "screenshot_som": np.full((4, 4, 3), 200, dtype=np.uint8), + } + if self.goal_object: + obs["goal_object"] = [{"type": "text", "text": "reach the end"}] + return obs + + def reset(self, seed=None, options=None): + self.reset_seed = seed + self.t = 0 + return self._obs(), {"recording_start_time": 123.0} + + def step(self, action): + self.actions.append(action) + self.t += 1 + terminated = self.t >= self.n_steps + info = { + "action_exec_start": 10.0, + "action_exec_stop": 13.0, + "action_exec_timeout": 1.0, + "RAW_REWARD_GLOBAL": 0.25 * self.t, + "task_info": {"t": self.t}, + } + return self._obs(), float(terminated), terminated, False, info + + def close(self): + self.closed = True + + +@pytest.fixture(autouse=True) +def restore_logging(): + root = logging.getLogger() + handlers, level = list(root.handlers), root.level + openai_logger = logging.getLogger("openai._base_client") + openai_level = openai_logger.level + yield + for h in list(root.handlers): + if h not in handlers: + root.removeHandler(h) + h.close() + root.handlers[:] = handlers + root.setLevel(level) + openai_logger.setLevel(openai_level) + + +@pytest.fixture +def no_debug(monkeypatch): + monkeypatch.setattr(loop, "_is_debugging", lambda: False) + + +def _make_exp(tmp_path, monkeypatch, env=None, **agent_kwargs): + env = env if env is not None else FakeEnv() + calls = [] + + def fake_make_env(self, action_mapping, exp_dir, exp_task_kwargs={}, use_raw_page_output=True): + calls.append(dict(action_mapping=action_mapping, exp_dir=exp_dir, raw=use_raw_page_output)) + return env + + monkeypatch.setattr(EnvArgs, "make_env", fake_make_env) + exp = ExpArgs( + agent_args=FakeAgentArgs(**agent_kwargs), + env_args=EnvArgs(task_name="fake.task", task_seed=7), + ) + exp.prepare(tmp_path) + return exp, env, calls + + +def _summary(exp): + return json.loads((Path(exp.exp_dir) / "summary_info.json").read_text()) + + +# --------------------------------------------------------------------------- EnvArgs.make_env + + +@pytest.fixture +def fake_gym_make(monkeypatch): + make = mock.MagicMock(return_value="ENV") + monkeypatch.setattr(loop.gym, "make", make) + monkeypatch.setattr(loop, "_get_env_name", lambda name: f"browsergym/{name}") + return make + + +def test_make_env_minimal_kwargs(fake_gym_make, tmp_path): + env = EnvArgs(task_name="miniwob.click-test", max_steps=3).make_env( + action_mapping="MAP", exp_dir=tmp_path + ) + assert env == "ENV" + fake_gym_make.assert_called_once_with( + "browsergym/miniwob.click-test", + disable_env_checker=True, + max_episode_steps=3, + headless=True, + wait_for_user_message=False, + action_mapping="MAP", + use_raw_page_output=True, + ) + + +def test_make_env_all_optional_kwargs(fake_gym_make, tmp_path): + args = EnvArgs( + task_name="miniwob.click-test", + record_video=True, + viewport={"width": 10, "height": 20}, + slow_mo=0, + pre_observation_delay=0.5, + storage_state={"cookies": []}, + task_kwargs={"a": 1, "b": 2}, + headless=False, + wait_for_user_message=True, + ) + args.make_env( + action_mapping=None, + exp_dir=tmp_path, + exp_task_kwargs={"b": 3, "c": 4}, + use_raw_page_output=False, + ) + kwargs = fake_gym_make.call_args.kwargs + assert kwargs["record_video_dir"] == tmp_path + assert kwargs["viewport"] == {"width": 10, "height": 20} + assert kwargs["slow_mo"] == 0 # 0 is not None -> forwarded + assert kwargs["pre_observation_delay"] == 0.5 + assert kwargs["pw_context_kwargs"] == {"storage_state": {"cookies": []}} + assert kwargs["task_kwargs"] == {"a": 1, "b": 3, "c": 4} + assert kwargs["headless"] is False + assert kwargs["wait_for_user_message"] is True + assert kwargs["use_raw_page_output"] is False + + +def test_make_env_exp_task_kwargs_without_task_kwargs(fake_gym_make, tmp_path): + EnvArgs(task_name="x").make_env(None, tmp_path, exp_task_kwargs={"z": 1}) + assert fake_gym_make.call_args.kwargs["task_kwargs"] == {"z": 1} + + +def test_make_env_assistantbench_output_file(fake_gym_make, tmp_path): + EnvArgs(task_name="assistantbench.test.3", task_kwargs={"k": "v"}).make_env(None, tmp_path) + assert fake_gym_make.call_args.kwargs["task_kwargs"] == { + "k": "v", + "output_file": tmp_path / "assistantbench-prediction.json", + } + + +def test_make_env_assistantbench_validation_is_not_patched(fake_gym_make, tmp_path): + EnvArgs(task_name="assistantbench.validation.3").make_env(None, tmp_path) + assert "task_kwargs" not in fake_gym_make.call_args.kwargs + + +@pytest.mark.parametrize( + "task_name, module", + [ + ("miniwob.click-test", "browsergym.miniwob"), + ("workarena.servicenow.x", "browsergym.workarena"), + ("webarena.1", "browsergym.webarena"), + ("visualwebarena.1", "browsergym.visualwebarena"), + ("assistantbench.test.1", "browsergym.assistantbench"), + ("weblinx.abc", "weblinx_browsergym"), + ("openended", None), + ], +) +def test_get_env_name_lazy_imports(monkeypatch, task_name, module): + if module is not None: + monkeypatch.setitem(sys.modules, module, types.ModuleType(module)) + assert loop._get_env_name(task_name) == f"browsergym/{task_name}" + + +# --------------------------------------------------------------------------- small helpers + + +def test_seed_max_is_xor_today(): + # Documents current behavior; see the xfail test below for the intended value. + assert loop.SEED_MAX == 34 + + +@pytest.mark.xfail(strict=True, reason="BUG: SEED_MAX = 2 ^ 32 is XOR (== 34), not 2**32") +def test_seed_max_is_two_to_the_32(): + assert loop.SEED_MAX == 2**32 + + +def test_abstract_agent_args_name_defaults_to_class_name(): + assert FakeAgentArgs().agent_name == "FakeAgentArgs" + assert FakeAgentArgs(agent_name="custom").agent_name == "custom" + args = FakeAgentArgs() + assert args.prepare() is None + assert args.close() is None + + +def test_abstract_agent_args_is_abstract(): + with pytest.raises(TypeError): + AbstractAgentArgs() + + +def test_save_package_versions(tmp_path): + loop.save_package_versions(tmp_path) + lines = (tmp_path / "package_versions.txt").read_text().splitlines() + assert lines == sorted(lines) + assert any(line.startswith("numpy==") for line in lines) + + +def test_extract_err_msg_returns_last_error(): + steps = [ + StepInfo(agent_info=None), + StepInfo(agent_info={"err_msg": "first", "stack_trace": "st1"}), + StepInfo(agent_info={}), + StepInfo(agent_info={"err_msg": "second"}), + ] + assert loop._extract_err_msg(steps) == ("second", None) + assert loop._extract_err_msg([StepInfo()]) == (None, None) + assert loop._extract_err_msg([]) == (None, None) + + +def test_aggregate_episode_stats(): + steps = [ + StepInfo(stats=None), + StepInfo(stats={"cost": 1.0, "tokens": np.int64(5), "gone": None}), + StepInfo(stats={"cost": 2.5, "tokens": np.int64(7), "gone": None}), + ] + with pytest.warns(RuntimeWarning): # nanmax over all-NaN column + agg = loop._aggregate_episode_stats(steps) + assert agg["cum_steps"] == 3 + assert agg["cum_cost"] == 3.5 and agg["max_cost"] == 2.5 + assert agg["cum_tokens"] == 12 and agg["max_tokens"] == 7 + assert type(agg["cum_tokens"]) is int + assert agg["cum_gone"] == 0.0 # nansum of all-NaN is 0 + assert agg["max_gone"] is None # nanmax of all-NaN is NaN -> None + json.dumps(agg) # must be JSON serialisable + + +def test_aggregate_episode_stats_empty(): + assert loop._aggregate_episode_stats([]) == {"cum_steps": 0} + + +def test_is_debugging_reflects_trace(monkeypatch): + monkeypatch.setattr(loop.sys, "gettrace", lambda: None) + assert loop._is_debugging() is False + monkeypatch.setattr(loop.sys, "gettrace", lambda: object()) + assert loop._is_debugging() is True + + +def test_flatten_dict(): + assert loop._flatten_dict({"a": 1, "b": {"c": 2, "d": {"e": 3}}, "f": {}}) == { + "a": 1, + "b.c": 2, + "b.d.e": 3, + } + assert loop._flatten_dict({"a": {"b": 1}}, sep="/") == {"a/b": 1} + + +@pytest.mark.parametrize( + "agent_info, expected", + [ + ({"think": "I should click"}, "I should click\n\naction:\nclick('a')\n"), + ({}, "action:\nclick('a')\n"), + ], +) +def test_send_chat_info(agent_info, expected): + chat = mock.MagicMock() + loop._send_chat_info(chat, "click('a')", agent_info) + chat.add_message.assert_called_once_with(role="info", msg=expected) + + +def test_move_old_exp(tmp_path): + exp_dir = tmp_path / "exp" + loop._move_old_exp(exp_dir) # no-op when missing + assert not (tmp_path / "_exp").exists() + exp_dir.mkdir() + (exp_dir / "f.txt").write_text("x") + loop._move_old_exp(str(exp_dir)) + assert not exp_dir.exists() + assert (tmp_path / "_exp" / "f.txt").read_text() == "x" + + +def test_dataclass_json_encoder(): + payload = { + "ts": StepTimestamps(env_start=1.5), + "i": np.int32(3), + "f": np.float32(0.5), + "arr": np.array([[1, 2]]), + } + decoded = json.loads(json.dumps(payload, cls=DataclassJSONEncoder)) + assert decoded["ts"]["env_start"] == 1.5 + assert decoded["i"] == 3 and decoded["f"] == 0.5 and decoded["arr"] == [[1, 2]] + with pytest.raises(TypeError): + json.dumps({"s": {1, 2}}, cls=DataclassJSONEncoder) + + +# --------------------------------------------------------------------------- StepInfo + + +def test_step_info_from_reset_and_step_profiling(): + env = FakeEnv(n_steps=1) + pre = mock.MagicMock(side_effect=lambda o: {**o, "pre": 1}) + + s0 = StepInfo(step=0) + s0.from_reset(env, seed=3, obs_preprocessor=pre) + assert env.reset_seed == 3 + assert (s0.reward, s0.terminated, s0.truncated) == (0, False, False) + assert s0.obs["pre"] == 1 + assert s0.profiling.action_exec_start == 123.0 + assert s0.profiling.action_exec_stop == s0.profiling.env_stop + assert not s0.is_done + + s1 = StepInfo(step=1) + s1.from_step(env, "act", obs_preprocessor=None) + assert env.actions == ["act"] + assert "pre" not in s1.obs + assert s1.reward == 1.0 and s1.terminated is True and s1.is_done + assert s1.raw_reward == 0.25 + assert s1.task_info == {"t": 1} + t = s1.profiling + assert t.action_exec_start == 10.0 + assert t.action_exect_after_timeout == 13.0 + assert t.action_exec_stop == 12.0 + assert t.wait_for_page_loading_start is None and t.validation_stop is None + + +def test_step_info_from_reset_without_recording_start_time(): + env = FakeEnv() + env.reset = lambda seed=None: ({"goal": "g"}, {}) + s = StepInfo(step=0) + s.from_reset(env, seed=0, obs_preprocessor=None) + assert s.profiling.action_exec_start == s.profiling.env_start + + +@pytest.mark.parametrize( + "terminated, truncated, done", + [(False, False, False), (True, False, True), (False, True, True), (None, None, None)], +) +def test_step_info_is_done(terminated, truncated, done): + assert StepInfo(terminated=terminated, truncated=truncated).is_done == done + + +def test_step_info_from_action_builds_stats(): + agent = FakeAgent(["click('x')"]) + s = StepInfo(step=0, obs={"goal": "hello world", "n": 3}) + assert s.from_action(agent) == "click('x')" + assert s.action == "click('x')" + assert "stats" not in s.agent_info # popped into s.stats + assert s.stats["llm_cost"] == 0.5 + assert s.stats["n_token_goal"] > 0 + assert "n_token_n" not in s.stats + assert s.stats["agent_elapsed"] >= 0 + assert "step_elapsed" in s.stats + + +def test_step_info_make_stats_non_dict_obs(): + s = StepInfo(obs="raw text", agent_info={}) + s.make_stats() + assert set(s.stats) == {"step_elapsed", "agent_elapsed"} + + +def test_save_step_info_writes_files_and_restores_screenshots(tmp_path): + shot = np.zeros((3, 3, 3), dtype=np.uint8) + som = np.full((3, 3, 3), 255, dtype=np.uint8) + s = StepInfo( + step=2, + obs={"goal": "g", "screenshot": shot, "screenshot_som": som, "goal_object": ["G1"]}, + agent_info={"think": "t"}, + ) + s.save_step_info(tmp_path, save_json=True, save_screenshot=True, save_som=True) + + assert (tmp_path / "screenshot_step_2.png").exists() + assert (tmp_path / "screenshot_som_step_2.png").exists() + with gzip.open(tmp_path / "goal_object.pkl.gz", "rb") as f: + assert pickle.load(f) == ["G1"] + with gzip.open(tmp_path / "step_2.pkl.gz", "rb") as f: + saved = pickle.load(f) + assert "screenshot" not in saved.obs and "screenshot_som" not in saved.obs + assert saved.obs["goal_object"] is None + json_info = json.loads((tmp_path / "steps_info.json").read_text()) + assert json_info["step"] == 2 and json_info["agent_info"] == {"think": "t"} + + # screenshots are put back on the live object; goal_object stays a placeholder + assert s.obs["screenshot"] is shot and s.obs["screenshot_som"] is som + assert s.obs["goal_object"] is None + + +def test_save_step_info_goal_object_written_once(tmp_path): + StepInfo(step=0, obs={"goal_object": ["first"]}).save_step_info(tmp_path) + StepInfo(step=1, obs={"goal_object": ["second"]}).save_step_info(tmp_path) + with gzip.open(tmp_path / "goal_object.pkl.gz", "rb") as f: + assert pickle.load(f) == ["first"] + + +def test_save_step_info_respects_flags(tmp_path): + s = StepInfo( + step=0, + obs={"screenshot": np.zeros((2, 2, 3), np.uint8), "screenshot_som": None}, + ) + s.save_step_info(tmp_path, save_screenshot=False, save_som=True) + assert not list(tmp_path.glob("screenshot*")) + assert not (tmp_path / "steps_info.json").exists() + assert "screenshot_som" not in s.obs # None is not restored + + +def test_save_step_info_non_dict_obs(tmp_path): + StepInfo(step=5, obs=None).save_step_info(tmp_path) + with gzip.open(tmp_path / "step_5.pkl.gz", "rb") as f: + assert pickle.load(f).obs is None + + +# --------------------------------------------------------------------------- ExpArgs.prepare / _make_dir + + +def test_make_id_is_stable(): + exp = ExpArgs(agent_args=FakeAgentArgs(), env_args=EnvArgs(task_name="t")) + exp.make_id() + first = exp.exp_id + assert len(first) == 36 + exp.make_id() + assert exp.exp_id == first + exp2 = ExpArgs(agent_args=None, env_args=None, exp_id="fixed") + exp2.make_id() + assert exp2.exp_id == "fixed" + + +def test_prepare_creates_dir_and_pickles(tmp_path): + exp = ExpArgs(agent_args=FakeAgentArgs(), env_args=EnvArgs(task_name="t.a/b:c", task_seed=None)) + with mock.patch.object(loop.np.random, "randint", return_value=11) as randint: + exp.prepare(tmp_path) + randint.assert_called_once_with(0, loop.SEED_MAX) + assert exp.env_args.task_seed == 11 + assert exp.exp_name == "FakeAgentArgs_on_t.a/b:c_11" + assert exp.exp_dir.parent == tmp_path + assert exp.exp_dir.name.endswith("_FakeAgentArgs_on_t.a_b_c_11") # sanitised + assert exp.exp_id is not None + with open(exp.exp_dir / "exp_args.pkl", "rb") as f: + reloaded = pickle.load(f) + assert reloaded.exp_name == exp.exp_name and reloaded.exp_id == exp.exp_id + + +def test_prepare_keeps_explicit_name_and_moves_old_dir(tmp_path): + old = tmp_path / "old_exp" + old.mkdir() + exp = ExpArgs( + agent_args=FakeAgentArgs(), + env_args=EnvArgs(task_name="t", task_seed=0), + exp_name="my exp", + exp_dir=old, + ) + exp.prepare(tmp_path / "root") + assert (tmp_path / "_old_exp").exists() and not old.exists() + assert exp.exp_dir.name.endswith("_my exp") + assert exp.env_args.task_seed == 0 # 0 is a valid seed and must not be resampled + + +def test_make_dir_adds_suffix_on_collision(tmp_path): + from datetime import datetime + + exp = ExpArgs(agent_args=None, env_args=None, exp_name="e") + exp.exp_date = datetime(2024, 1, 2, 3, 4, 5) + (tmp_path / "2024-01-02_03-04-05_e").mkdir() + (tmp_path / "2024-01-02_03-04-05_e_1").mkdir() + exp._make_dir(tmp_path) + assert exp.exp_dir == tmp_path / "2024-01-02_03-04-05_e_2" + + +def test_make_dir_gives_up_after_999(tmp_path): + from datetime import datetime + + exp = ExpArgs(agent_args=None, env_args=None, exp_name="e") + exp.exp_date = datetime(2024, 1, 2) + with mock.patch.object(loop.Path, "exists", return_value=True): + with pytest.raises(ValueError, match="unique name"): + exp._make_dir(tmp_path) + + +# --------------------------------------------------------------------------- ExpArgs.run + + +def test_run_happy_path(tmp_path, monkeypatch, no_debug): + env = FakeEnv(n_steps=2, with_chat=True) + exp, env, calls = _make_exp(tmp_path, monkeypatch, env=env) + exp.run() + exp_dir = Path(exp.exp_dir) + + agent = exp.agent_args.agents[0] + assert agent.task_name == "fake.task" + assert calls[0]["exp_dir"] == exp_dir and calls[0]["raw"] is False + assert calls[0]["action_mapping"]("x") == "# x" + assert env.reset_seed == 7 + assert env.actions == ["click('a')", "click('b')"] + assert env.closed + + summary = _summary(exp) + assert summary["n_steps"] == 2 + assert summary["cum_reward"] == 1.0 + assert summary["cum_raw_reward"] == 0.75 + assert summary["err_msg"] is None and summary["stack_trace"] is None + assert summary["terminated"] is True and summary["truncated"] is False + assert summary["stats.cum_steps"] == 3 + assert summary["stats.cum_llm_cost"] == 1.0 + + for i in range(3): + assert (exp_dir / f"step_{i}.pkl.gz").exists() + assert (exp_dir / f"screenshot_step_{i}.png").exists() + assert not list(exp_dir.glob("screenshot_som_*")) # save_som defaults to False + assert (exp_dir / "goal_object.pkl.gz").exists() + assert (exp_dir / "package_versions.txt").exists() + assert "Running experiment" in (exp_dir / "experiment.log").read_text() + assert env.chat.add_message.call_count == 2 + assert exp.err_msg is None + + # the experiment file handler is detached after run + root = logging.getLogger() + assert exp.logging_file_handler not in root.handlers + + result = ExpResult(exp_dir) + assert result.status == "done" + steps = result.steps_info + assert [s.step for s in steps] == [0, 1, 2] + assert steps[0].action == "click('a')" and steps[2].action is None + assert steps[1].obs["preprocessed"] is True + assert steps[1].obs["goal_object"] == [{"type": "text", "text": "reach the end"}] + assert steps[1].obs["screenshot"].shape == (4, 4, 3) + assert steps[1].obs["screenshot"][0, 0, 0] == 10 + assert "screenshot_som" not in steps[1].obs + assert result.exp_args.exp_id == exp.exp_id + + +def test_run_none_action_truncates(tmp_path, monkeypatch, no_debug): + exp, env, _ = _make_exp(tmp_path, monkeypatch, env=FakeEnv(n_steps=5), actions=()) + exp.save_som = True + exp.run() + summary = _summary(exp) + assert summary["n_steps"] == 0 + assert summary["truncated"] is True and summary["terminated"] is False + assert summary["err_msg"] is None + assert env.actions == [] + assert (Path(exp.exp_dir) / "screenshot_som_step_0.png").exists() + assert ExpResult(exp.exp_dir).status == "done" + + +def test_run_agent_exception_is_captured(tmp_path, monkeypatch, no_debug): + exp, env, _ = _make_exp(tmp_path, monkeypatch, raise_at=1) + exp.run() + summary = _summary(exp) + assert summary["err_msg"].startswith( + "Exception uncaught by agent or environment in task fake.task.\nValueError:\nagent exploded" + ) + assert "Traceback" in summary["stack_trace"] + assert summary["n_steps"] == 1 + assert summary["terminated"] is False + assert exp.err_msg == summary["err_msg"] and exp.stack_trace == summary["stack_trace"] + assert env.closed + # the failing step is still saved by the finally block + assert (Path(exp.exp_dir) / "step_1.pkl.gz").exists() + assert ExpResult(exp.exp_dir).status == "error" + + +def test_run_make_env_failure(tmp_path, monkeypatch, no_debug): + exp, _, _ = _make_exp(tmp_path, monkeypatch) + + def boom(self, **kwargs): + raise RuntimeError("no browser") + + monkeypatch.setattr(EnvArgs, "make_env", boom) + exp.run() + summary = _summary(exp) + assert "RuntimeError:\nno browser" in summary["err_msg"] + assert summary["n_steps"] == 0 + assert summary["terminated"] is None and summary["truncated"] is None + assert not list(Path(exp.exp_dir).glob("step_*.pkl.gz")) + + +def test_run_agent_error_in_agent_info_is_promoted(tmp_path, monkeypatch, no_debug): + exp, _, _ = _make_exp(tmp_path, monkeypatch) + orig = FakeAgent.get_action + + def get_action(self, obs): + action, info = orig(self, obs) + if self.calls == 1: + info["err_msg"] = "parse error" + info["stack_trace"] = "st" + return action, info + + monkeypatch.setattr(FakeAgent, "get_action", get_action) + exp.run() + summary = _summary(exp) + assert summary["err_msg"] == "parse error" and summary["stack_trace"] == "st" + + +def test_run_reraises_in_debug_mode(tmp_path, monkeypatch): + monkeypatch.setattr(loop, "_is_debugging", lambda: True) + exp, env, _ = _make_exp(tmp_path, monkeypatch, raise_at=0) + with pytest.raises(ValueError, match="agent exploded"): + exp.run() + assert "ValueError" in _summary(exp)["err_msg"] # finally block still ran + assert env.closed + + +def test_run_debug_mode_disabled_by_flag(tmp_path, monkeypatch): + monkeypatch.setattr(loop, "_is_debugging", lambda: True) + exp, _, _ = _make_exp(tmp_path, monkeypatch, raise_at=0) + exp.enable_debug = False + exp.run() + assert "ValueError" in _summary(exp)["err_msg"] + + +def test_run_keyboard_interrupt_marks_early_termination(tmp_path, monkeypatch, no_debug): + exp, env, _ = _make_exp(tmp_path, monkeypatch) + monkeypatch.setattr(FakeAgent, "get_action", mock.Mock(side_effect=KeyboardInterrupt)) + with pytest.raises(KeyboardInterrupt): + exp.run() + summary = _summary(exp) + assert "KeyboardInterrupt:\nEarly termination??" in summary["err_msg"] + assert summary["stack_trace"] is None + assert env.closed + + +def test_run_finally_block_swallows_cleanup_errors(tmp_path, monkeypatch, no_debug, caplog): + exp, env, _ = _make_exp(tmp_path, monkeypatch, env=FakeEnv(n_steps=1)) + env.close = mock.Mock(side_effect=RuntimeError("close failed")) + monkeypatch.setattr(ExpArgs, "save_summary_info", mock.Mock(side_effect=OSError("disk"))) + monkeypatch.setattr(ExpArgs, "_unset_logger", mock.Mock(side_effect=RuntimeError("unset"))) + exp.run() # must not raise + env.close.assert_called_once() + ExpArgs._unset_logger.assert_called_once() + log = (Path(exp.exp_dir) / "experiment.log").read_text() + assert "Error while saving experiment info: disk" in log + assert "Error while closing the environment: close failed" in log + + +def test_run_finally_block_survives_save_step_failure(tmp_path, monkeypatch, no_debug): + exp, _, _ = _make_exp(tmp_path, monkeypatch, raise_at=0) + monkeypatch.setattr(StepInfo, "save_step_info", mock.Mock(side_effect=OSError("nope"))) + exp.run() + assert ( + "Error while saving step info in the finally block: nope" + in (Path(exp.exp_dir) / "experiment.log").read_text() + ) + assert "ValueError" in _summary(exp)["err_msg"] + + +def test_run_saves_tape_for_tape_agents(tmp_path, monkeypatch, no_debug): + save_tape = mock.Mock() + monkeypatch.setattr(loop, "TapeAgent", FakeAgent) + monkeypatch.setattr(loop, "save_tape", save_tape, raising=False) + env = FakeEnv(n_steps=1) + env.task = {"id": 1} + exp, _, _ = _make_exp(tmp_path, monkeypatch, env=env) + exp.run() + save_tape.assert_called_once() + exp_dir, episode_info, task, tape = save_tape.call_args.args + assert exp_dir == exp.exp_dir and task == {"id": 1} and tape == "the-tape" + assert [s.step for s in episode_info] == [0, 1] + + +def test_set_logger_configures_root_and_openai(tmp_path): + exp = ExpArgs(agent_args=None, env_args=None, logging_level=logging.DEBUG) + exp.exp_dir = tmp_path + exp._set_logger() + root = logging.getLogger() + assert exp.logging_file_handler in root.handlers + assert root.level == logging.DEBUG + assert logging.getLogger("openai._base_client").level == logging.INFO + exp._unset_logger() + assert exp.logging_file_handler not in root.handlers + exp.logging_file_handler.close() + + +@pytest.mark.xfail( + strict=True, + reason="BUG: _set_logger removes handlers while iterating root.handlers, skipping every other one", +) +def test_set_logger_removes_all_previous_stream_handlers(tmp_path): + root = logging.getLogger() + h1, h2 = logging.StreamHandler(), logging.StreamHandler() + root.handlers[:] = [h1, h2] + exp = ExpArgs(agent_args=None, env_args=None) + exp.exp_dir = tmp_path + exp._set_logger() + try: + assert h1 not in root.handlers and h2 not in root.handlers + finally: + exp.logging_file_handler.close() + + +# --------------------------------------------------------------------------- ExpResult + + +def _write_summary(exp_dir, **summary): + (Path(exp_dir) / "summary_info.json").write_text(json.dumps(summary)) + + +@pytest.mark.parametrize( + "summary, status", + [ + (None, "incomplete"), + ({"err_msg": "boom", "terminated": True}, "error"), + ({"err_msg": None, "terminated": True}, "done"), + ({"truncated": True}, "done"), + ({"terminated": False, "truncated": False}, "incomplete"), + ], +) +def test_exp_result_status(tmp_path, summary, status): + if summary is not None: + _write_summary(tmp_path, **summary) + assert ExpResult(tmp_path).status == status + + +def test_exp_result_empty_summary_is_incomplete(tmp_path): + (tmp_path / "summary_info.json").write_text("") + res = ExpResult(tmp_path) + with pytest.raises(FileNotFoundError, match="empty"): + res.summary_info + assert res.status == "incomplete" + + +def test_exp_result_summary_is_cached(tmp_path): + _write_summary(tmp_path, a=1) + res = ExpResult(tmp_path) + assert res.summary_info == {"a": 1} + _write_summary(tmp_path, a=2) + assert res.summary_info == {"a": 1} + + +def test_exp_result_exp_args_follows_moved_dir(tmp_path): + exp = ExpArgs(agent_args=FakeAgentArgs(), env_args=EnvArgs(task_name="t", task_seed=1)) + exp.prepare(tmp_path / "a") + moved = tmp_path / "moved" + Path(exp.exp_dir).rename(moved) + res = ExpResult(str(moved)) + assert res.exp_args.exp_dir == moved + assert res.exp_args is res.exp_args + flat = res.flat_exp_args + assert flat["env_args.task_name"] == "t" + assert flat["agent_args.agent_name"] == "FakeAgentArgs" + assert res.flat_exp_args is flat + + +def test_exp_result_get_exp_record(tmp_path): + res = ExpResult(tmp_path) + assert res.get_exp_record() == {"exp_dir": tmp_path} + + exp = ExpArgs(agent_args=FakeAgentArgs(), env_args=EnvArgs(task_name="t", task_seed=1)) + exp.prepare(tmp_path / "r") + _write_summary(exp.exp_dir, cum_reward=1.0) + record = ExpResult(exp.exp_dir).get_exp_record() + assert record["cum_reward"] == 1.0 and record["env_args.task_seed"] == 1 + + +def test_exp_result_screenshots(tmp_path): + Image.new("RGB", (2, 2), (1, 2, 3)).save(tmp_path / "screenshot_step_0.png") + Image.new("L", (2, 2), 9).save(tmp_path / "screenshot_step_2.jpg") + Image.new("RGB", (2, 2), (5, 5, 5)).save(tmp_path / "screenshot_som_step_0.png") + res = ExpResult(tmp_path) + + assert res.get_screenshot_path(0) == tmp_path / "screenshot_step_0.png" + assert res.get_screenshot_path(2) == tmp_path / "screenshot_step_2.jpg" + with pytest.raises(FileNotFoundError, match="step 1"): + res.get_screenshot_path(1) + + shots = res.screenshots + assert len(shots) == 3 and shots[1] is None + assert shots[0].getpixel((0, 0)) == (1, 2, 3) + assert shots[2].mode == "RGB" + assert res.get_screenshot(0) is shots[0] # cached + + with pytest.raises(FileNotFoundError): # som listing is driven by regular screenshots + res.screenshots_som + + +def test_exp_result_step_info_without_files(tmp_path): + StepInfo(step=0, obs={"goal": "g"}).save_step_info(tmp_path) + StepInfo(step=1, obs=None).save_step_info(tmp_path) + res = ExpResult(tmp_path) + s0 = res.get_step_info(0) + assert "screenshot" not in s0.obs and "screenshot_som" not in s0.obs + assert res.get_step_info(0) is s0 + assert res.get_step_info(1).obs is None + assert [s.step for s in res.steps_info] == [0, 1] + + +def test_exp_result_videos_and_logs(tmp_path): + res = ExpResult(tmp_path) + with pytest.raises(FileNotFoundError, match="chat_video"): + res.chat_video_path + with pytest.raises(FileNotFoundError, match="task_video"): + res.task_video_path + (tmp_path / "chat_video").mkdir() + (tmp_path / "chat_video" / "c.webm").write_bytes(b"") + (tmp_path / "task_video").mkdir() + (tmp_path / "task_video" / "t.webm").write_bytes(b"") + assert res.chat_video_path.name == "c.webm" + assert res.task_video_path.name == "t.webm" + assert res.combined_video_path == tmp_path / "combined_video.mp4" + (tmp_path / "experiment.log").write_text("hello") + assert res.logs == "hello" + (tmp_path / "experiment.log").write_text("changed") + assert res.logs == "hello" + + +# --------------------------------------------------------------------------- discovery helpers + + +@pytest.fixture +def empty_cache(monkeypatch): + cache = {} + monkeypatch.setattr(loop, "EXP_RESULT_CACHE", cache) + return cache + + +def test_get_exp_result_caches_by_str_path(tmp_path, empty_cache): + a = loop.get_exp_result(tmp_path) + b = loop.get_exp_result(str(tmp_path)) + assert a is b and list(empty_cache) == [str(tmp_path)] + + +def _touch_exp(path): + path.mkdir(parents=True) + (path / "exp_args.pkl").write_bytes(b"") + return path + + +def test_yield_all_exp_results(tmp_path, empty_cache): + root1, root2 = tmp_path / "r1", tmp_path / "r2" + visible = _touch_exp(root1 / "study" / "exp_a") + hidden_u = _touch_exp(root1 / "_exp_b") + hidden_d = _touch_exp(root2 / ".exp_c") + other = _touch_exp(root2 / "exp_d") + + found = {r.exp_dir for r in loop.yield_all_exp_results([root1, root2], progress_fn=None)} + assert found == {visible, other} + + progress = mock.Mock(side_effect=lambda it, desc: it) + found = { + r.exp_dir + for r in loop.yield_all_exp_results( + [root1, root2], progress_fn=progress, load_hidden=True, use_cache=False + ) + } + assert found == {visible, other, hidden_u, hidden_d} + progress.assert_called_once() + assert len(empty_cache) == 2 # only the cached first pass populated it + + single = list(loop.yield_all_exp_results(str(root2), progress_fn=None)) + assert [r.exp_dir for r in single] == [other] diff --git a/tests/test_ui_assistant_unit.py b/tests/test_ui_assistant_unit.py new file mode 100644 index 000000000..f597678fc --- /dev/null +++ b/tests/test_ui_assistant_unit.py @@ -0,0 +1,121 @@ +"""Hermetic unit tests for agentlab.ui_assistant (no browser, no LLM).""" + +import copy +import sys +from dataclasses import dataclass +from types import SimpleNamespace +from unittest import mock + +import pytest + +from agentlab import ui_assistant +from agentlab.agents.generic_agent import AGENT_4o +from agentlab.experiments.loop import AbstractAgentArgs, EnvArgs, ExpArgs +from agentlab.ui_assistant import make_exp_args + + +@dataclass +class PlainAgentArgs(AbstractAgentArgs): + def make_agent(self): + raise NotImplementedError + + +def test_make_exp_args_generic_agent(): + agent_args = copy.deepcopy(AGENT_4o) + agent_args.flags.action.demo_mode = "off" + agent_args.flags.enable_chat = False + + exp_args = make_exp_args(agent_args, "https://example.test") + + assert isinstance(exp_args, ExpArgs) + assert exp_args.agent_args is agent_args + assert agent_args.flags.action.demo_mode == "default" + assert agent_args.flags.enable_chat is True + env = exp_args.env_args + assert isinstance(env, EnvArgs) + assert env.task_name == "openended" + assert env.task_kwargs == {"start_url": "https://example.test"} + assert env.task_seed is None + assert env.max_steps == 1000 + assert env.headless is False + assert env.record_video is True + assert env.wait_for_user_message is True + assert env.viewport == {"width": 1500, "height": 1280} + assert env.slow_mo == 1000 + assert exp_args.exp_dir is None and exp_args.exp_name is None + + +def test_make_exp_args_non_generic_agent_with_flags(): + agent_args = SimpleNamespace(flags=SimpleNamespace(action=SimpleNamespace(demo_mode="off"))) + exp_args = make_exp_args(agent_args, "https://a.test") + assert agent_args.flags.action.demo_mode == "default" + assert not hasattr(agent_args.flags, "enable_chat") # only GenericAgentArgs gets chat + assert exp_args.env_args.task_kwargs == {"start_url": "https://a.test"} + + +def test_make_exp_args_agent_without_flags(): + agent_args = PlainAgentArgs() + exp_args = make_exp_args(agent_args, "https://b.test") + assert exp_args.agent_args is agent_args + assert not hasattr(agent_args, "flags") + + +@pytest.fixture +def fake_main_env(monkeypatch, tmp_path): + agent_args = PlainAgentArgs() + import_object = mock.Mock(return_value=agent_args) + prepare, run = mock.Mock(), mock.Mock() + monkeypatch.setattr(ui_assistant, "import_object", import_object) + monkeypatch.setattr(ui_assistant, "RESULTS_DIR", tmp_path) + monkeypatch.setattr(ExpArgs, "prepare", prepare) + monkeypatch.setattr(ExpArgs, "run", run) + return SimpleNamespace( + agent_args=agent_args, import_object=import_object, prepare=prepare, run=run + ) + + +def test_main_defaults(monkeypatch, tmp_path, fake_main_env): + monkeypatch.setattr(sys, "argv", ["ui_assistant"]) + created = [] + real_make = ui_assistant.make_exp_args + monkeypatch.setattr( + ui_assistant, + "make_exp_args", + lambda a, url: created.append(real_make(a, url)) or created[-1], + ) + + ui_assistant.main() + + fake_main_env.import_object.assert_called_once_with( + "agentlab.agents.generic_agent.AGENT_GPT5_MINI" + ) + (exp_args,) = created + assert exp_args.agent_args is fake_main_env.agent_args + assert exp_args.env_args.task_kwargs == {"start_url": "https://www.google.com"} + fake_main_env.prepare.assert_called_once_with(tmp_path / "ui_assistant_logs") + fake_main_env.run.assert_called_once_with() + + +def test_main_custom_args_and_unknown_flags(monkeypatch, tmp_path, fake_main_env): + monkeypatch.setattr( + sys, + "argv", + [ + "ui_assistant", + "--agent_config", + "my.pkg.AGENT", + "--start_url", + "https://custom.test", + "--unknown", + "ignored", + ], + ) + make = mock.Mock(wraps=ui_assistant.make_exp_args) + monkeypatch.setattr(ui_assistant, "make_exp_args", make) + + ui_assistant.main() + + fake_main_env.import_object.assert_called_once_with("my.pkg.AGENT") + make.assert_called_once_with(fake_main_env.agent_args, "https://custom.test") + fake_main_env.prepare.assert_called_once() + fake_main_env.run.assert_called_once()