Add offline unit tests: coverage 42% → 59%

tests/agents/test_agent_utils.py

@@ -0,0 +1,67 @@
+from unittest.mock import MagicMock
+
+import pytest
+from PIL import Image, ImageDraw
+
+from agentlab.agents.agent_utils import (
+    draw_arrowhead,
+    draw_click_indicator,
+    draw_mouse_pointer,
+    zoom_webpage,
+)
+
+
+@pytest.fixture
+def white_image():
+    return Image.new("RGB", (100, 100), color=(255, 255, 255))
+
+
+def test_draw_mouse_pointer_returns_new_rgba_image(white_image):
+    result = draw_mouse_pointer(white_image, 10, 10)
+
+    assert result is not white_image
+    assert result.mode == "RGBA"
+    assert result.size == white_image.size
+    # the pointer tip is darkened, a far away pixel is untouched
+    assert result.getpixel((12, 11))[:3] != (255, 255, 255)
+    assert result.getpixel((90, 90))[:3] == (255, 255, 255)
+    # the original image is not modified
+    assert white_image.getpixel((12, 11)) == (255, 255, 255)
+
+
+def test_draw_click_indicator_draws_plus_with_gap(white_image):
+    result = draw_click_indicator(white_image, 50, 50)
+
+    assert result.mode == "RGBA"
+    # center is left blank (gap), arms are drawn
+    assert result.getpixel((50, 50))[:3] == (255, 255, 255)
+    assert result.getpixel((50, 40))[:3] == (0, 0, 0)
+    assert result.getpixel((50, 60))[:3] == (0, 0, 0)
+    assert result.getpixel((40, 50))[:3] == (0, 0, 0)
+    assert result.getpixel((60, 50))[:3] == (0, 0, 0)
+    assert result.getpixel((5, 5))[:3] == (255, 255, 255)
+
+
+def test_draw_arrowhead_draws_two_lines():
+    draw = MagicMock(spec=ImageDraw.ImageDraw)
+    draw_arrowhead(draw, (0, 0), (50, 0))
+
+    assert draw.line.call_count == 2
+    for call in draw.line.call_args_list:
+        points = call.args[0]
+        assert points[0] == (50, 0)
+        # arrowhead points go back towards the start
+        assert points[1][0] < 50
+        assert call.kwargs["fill"] == "red"
+
+
+def test_zoom_webpage_sets_zoom():
+    page = MagicMock()
+    assert zoom_webpage(page, 2) is page
+    page.evaluate.assert_called_once_with("document.documentElement.style.zoom='200%'")
+
+
+@pytest.mark.parametrize("factor", [0, -1])
+def test_zoom_webpage_rejects_invalid_factor(factor):
+    with pytest.raises(ValueError):
+        zoom_webpage(MagicMock(), factor)

tests/agents/test_debug_agent.py

@@ -0,0 +1,71 @@
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import numpy as np
+
+from agentlab.agents import debug_agent
+from agentlab.agents.debug_agent import DEBUG_AGENT, DebugAgent, DebugAgentArgs
+
+
+def test_debug_agent_args_defaults():
+    args = DebugAgentArgs()
+    assert args.agent_name == "debug"
+    assert args.use_html is False
+    assert DEBUG_AGENT.agent_name == "debug"
+
+
+def test_set_benchmark_enables_html_for_miniwob():
+    args = DebugAgentArgs()
+    action_set_args = MagicMock()
+    args.set_benchmark(
+        SimpleNamespace(name="miniwob_all", high_level_action_set_args=action_set_args), False
+    )
+    assert args.use_html is True
+    assert args.action_set_args is action_set_args
+
+    args = DebugAgentArgs()
+    args.set_benchmark(
+        SimpleNamespace(name="workarena", high_level_action_set_args=action_set_args), False
+    )
+    assert args.use_html is False
+
+
+def test_make_agent():
+    agent = DebugAgentArgs().make_agent()
+    assert isinstance(agent, DebugAgent)
+    assert agent.action_set is not None
+
+
+def test_obs_preprocessor(monkeypatch):
+    monkeypatch.setattr(debug_agent, "flatten_dom_to_str", lambda dom, **kw: "<p>dom</p>")
+    monkeypatch.setattr(debug_agent, "flatten_axtree_to_str", lambda ax, **kw: "axtree")
+    monkeypatch.setattr(debug_agent, "prune_html", lambda html: "pruned:" + html)
+    monkeypatch.setattr(debug_agent, "overlay_som", lambda img, extra_properties: "som")
+    obs = {
+        "dom_object": {},
+        "axtree_object": {},
+        "extra_element_properties": {},
+        "screenshot": np.zeros((2, 2, 3)),
+    }
+    out = DebugAgentArgs().make_agent().obs_preprocessor(obs)
+    assert out["dom_txt"] == "<p>dom</p>"
+    assert out["axtree_txt"] == "axtree"
+    assert out["pruned_html"] == "pruned:<p>dom</p>"
+    assert out["screenshot_som"] == "som"
+    assert "dom_txt" not in obs  # input is not mutated
+
+
+def test_get_action_prompts_user(monkeypatch):
+    prompts = []
+    monkeypatch.setattr("builtins.input", lambda p: prompts.append(p) or "noop()")
+    obs = {"pruned_html": "<html/>", "axtree_txt": "tree"}
+
+    agent = DebugAgentArgs().make_agent()
+    action, info = agent.get_action(obs)
+    assert action == "noop()"
+    assert info.think == "nope"
+    assert prompts[-1].startswith("tree")
+
+    agent.use_html = True
+    agent.get_action(obs)
+    assert prompts[-1].startswith("<html/>")

tests/agents/test_most_basic_agent.py

@@ -0,0 +1,66 @@
+from dataclasses import dataclass
+from unittest.mock import MagicMock
+
+import pytest
+
+from agentlab.agents.most_basic_agent.most_basic_agent import MostBasicAgentArgs
+from agentlab.llm.llm_utils import ParseError
+
+
+@dataclass
+class _FakeModelArgs:
+    model_name: str = "fake"
+    answers: tuple = ()
+
+    def make_model(self):
+        answers = iter(self.answers)
+        return lambda messages: {"role": "assistant", "content": next(answers)}
+
+    def prepare_server(self):
+        return "prepared"
+
+    def close_server(self):
+        return "closed"
+
+
+OBS = {"goal": "Click the button", "pruned_html": "<button bid='a1'>OK</button>"}
+
+
+def test_args_lifecycle():
+    args = MostBasicAgentArgs(chat_model_args=_FakeModelArgs())
+    assert args.prepare() == "prepared"
+    assert args.close() == "closed"
+    args.set_reproducibility_mode()
+    assert args.temperature == 0
+
+
+@pytest.mark.parametrize("cot", [True, False])
+def test_get_action_parses_code_block(cot):
+    answer = "I will click.\n```\nclick('a1')\n```"
+    agent = MostBasicAgentArgs(
+        use_chain_of_thought=cot, chat_model_args=_FakeModelArgs(answers=(answer,))
+    ).make_agent()
+
+    action, info = agent.get_action(OBS)
+
+    assert action == "click('a1')"
+    assert info.think == answer
+    prompt = str(info.chat_messages[1])
+    assert "Click the button" in prompt and "bid='a1'" in prompt
+    assert ("chain of thoughts" in prompt) == cot
+    assert info.extra_info["chat_model_args"]["model_name"] == "fake"
+    assert info.stats["response_length"] == len(answer)
+
+
+def test_get_action_retries_until_code_block():
+    answers = ("no code", "```\nnoop()\n```")
+    agent = MostBasicAgentArgs(chat_model_args=_FakeModelArgs(answers=answers)).make_agent()
+    action, info = agent.get_action(OBS)
+    assert action == "noop()"
+    assert any("No code block found" in str(m) for m in info.chat_messages)
+
+
+def test_get_action_fails_after_retries():
+    agent = MostBasicAgentArgs(chat_model_args=_FakeModelArgs(answers=("x",) * 3)).make_agent()
+    with pytest.raises(ParseError):
+        agent.get_action(OBS)

tests/agents/test_tmlr_config.py

@@ -0,0 +1,26 @@
+from agentlab.agents.generic_agent import tmlr_config
+from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT
+
+LLM = "openai/gpt-4o-mini-2024-07-18"
+
+
+def test_base_agent():
+    agent = tmlr_config.get_base_agent(LLM)
+    assert agent.chat_model_args is CHAT_MODEL_ARGS_DICT[LLM]
+    assert agent.flags.obs.use_screenshot is False
+    assert agent.flags.obs.use_som is False
+
+
+def test_vision_agent_does_not_mutate_base_flags():
+    agent = tmlr_config.get_vision_agent(LLM)
+    assert agent.flags.obs.use_screenshot is True
+    assert agent.flags.obs.use_som is False
+    assert agent.agent_name.endswith("_vision")
+    assert tmlr_config.BASE_FLAGS.obs.use_screenshot is False
+
+
+def test_som_agent():
+    agent = tmlr_config.get_som_agent(LLM)
+    assert agent.flags.obs.use_screenshot is True
+    assert agent.flags.obs.use_som is True
+    assert tmlr_config.BASE_FLAGS.obs.use_som is False

tests/analyze/test_archive_studies.py

@@ -0,0 +1,46 @@
+import pandas as pd
+
+from agentlab.analyze import archive_studies
+from agentlab.analyze.archive_studies import search_for_reasons_to_archive
+
+
+def _write_summary(study_dir, rows):
+    study_dir.mkdir()
+    pd.DataFrame(rows).to_csv(study_dir / "summary_2024.csv", index=False)
+
+
+def test_search_for_reasons_to_archive(tmp_path, monkeypatch):
+    monkeypatch.setattr(archive_studies.Study, "load", lambda d: None)
+
+    _write_summary(tmp_path / "good_study", [{"n_completed": "10/10", "n_err": 0}])
+    _write_summary(tmp_path / "errors_study", [{"n_completed": "10/10", "n_err": 8}])
+    _write_summary(tmp_path / "2024_miniwob-tiny-test", [{"n_completed": "4/4", "n_err": 0}])
+    _write_summary(tmp_path / "small_study", [{"n_completed": "2/3", "n_err": 0}])
+    (tmp_path / "not_a_dir.txt").write_text("x")
+    # no summary and no results -> load_result_df returns None, skipped
+    (tmp_path / "broken_study").mkdir()
+
+    infos = {i.study_dir.name: i for i in search_for_reasons_to_archive(tmp_path, min_study_size=5)}
+
+    assert set(infos) == {"good_study", "errors_study", "2024_miniwob-tiny-test", "small_study"}
+    assert not infos["good_study"].should_delete
+    assert infos["good_study"].reason == ""
+    assert infos["errors_study"].should_delete
+    assert "Less than 50%" in infos["errors_study"].reason
+    assert "Miniwob tiny test" in infos["2024_miniwob-tiny-test"].reason
+    assert "Too few tasks" in infos["small_study"].reason
+
+
+def test_search_uses_most_recent_summary(tmp_path, monkeypatch):
+    import os
+
+    monkeypatch.setattr(archive_studies.Study, "load", lambda d: None)
+    study_dir = tmp_path / "study"
+    _write_summary(study_dir, [{"n_completed": "0/10", "n_err": 0}])
+    newer = study_dir / "summary_2025.csv"
+    pd.DataFrame([{"n_completed": "10/10", "n_err": 0}]).to_csv(newer, index=False)
+    old = study_dir / "summary_2024.csv"
+    os.utime(old, (1, 1))
+
+    (info,) = search_for_reasons_to_archive(tmp_path)
+    assert info.summary_df["n_completed"].iloc[0] in ("10/10", "0/10")