Add unit tests for untested pure-logic modules (coverage 32% → 49%)

tests/agents/test_agent_utils.py

@@ -0,0 +1,61 @@
+from unittest.mock import MagicMock
+
+import pytest
+from PIL import Image
+
+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), (255, 255, 255))
+
+
+def test_draw_mouse_pointer_returns_new_rgba_image(white_image):
+    out = draw_mouse_pointer(white_image, 10, 10)
+    assert out is not white_image
+    assert out.mode == "RGBA"
+    assert out.size == white_image.size
+    # pixel inside the arrow is darkened, pixel far away is untouched
+    assert out.getpixel((15, 13))[:3] != (255, 255, 255)
+    assert out.getpixel((90, 90))[:3] == (255, 255, 255)
+    # original is not modified
+    assert white_image.getpixel((15, 13)) == (255, 255, 255)
+
+
+def test_draw_click_indicator_leaves_center_gap(white_image):
+    out = draw_click_indicator(white_image, 50, 50)
+    assert out.mode == "RGBA"
+    assert out.getpixel((50, 50))[:3] == (255, 255, 255)  # gap at the center
+    assert out.getpixel((50, 40))[:3] == (0, 0, 0)  # top line
+    assert out.getpixel((60, 50))[:3] == (0, 0, 0)  # right line
+    assert out.getpixel((5, 5))[:3] == (255, 255, 255)
+
+
+def test_draw_arrowhead_draws_two_lines():
+    draw = MagicMock()
+    draw_arrowhead(draw, (0, 0), (100, 0))
+    assert draw.line.call_count == 2
+    for call in draw.line.call_args_list:
+        points = call.args[0]
+        assert points[0] == (100, 0)
+        assert call.kwargs == {"fill": "red", "width": 4}
+
+
+def test_zoom_webpage_evaluates_js():
+    page = MagicMock()
+    assert zoom_webpage(page, 2.0) is page
+    page.evaluate.assert_called_once_with("document.documentElement.style.zoom='200.0%'")
+
+
+@pytest.mark.parametrize("factor", [0, -1.5])
+def test_zoom_webpage_invalid_factor(factor):
+    page = MagicMock()
+    with pytest.raises(ValueError):
+        zoom_webpage(page, factor)
+    page.evaluate.assert_not_called()

tests/agents/test_most_basic_agent.py

@@ -0,0 +1,103 @@
+from dataclasses import dataclass, field
+
+import pytest
+
+from agentlab.agents.most_basic_agent import most_basic_agent
+from agentlab.agents.most_basic_agent.most_basic_agent import MostBasicAgent, MostBasicAgentArgs
+from agentlab.llm.llm_utils import AIMessage, Discussion, ParseError
+
+OBS = {"goal": "Click the OK button", "pruned_html": "<button bid='a1'>OK</button>"}
+
+
+class FakeChat:
+    def __init__(self, responses):
+        self.responses = list(responses)
+        self.calls = []
+
+    def __call__(self, messages):
+        self.calls.append(messages)
+        return AIMessage(self.responses.pop(0))
+
+
+@dataclass
+class FakeModelArgs:
+    model_name: str = "fake-model"
+    responses: list = field(default_factory=list)
+    prepared: bool = False
+    closed: bool = False
+
+    def make_model(self):
+        self.chat = FakeChat(self.responses)
+        return self.chat
+
+    def prepare_server(self):
+        self.prepared = True
+        return "prepared"
+
+    def close_server(self):
+        self.closed = True
+        return "closed"
+
+
+def test_args_make_agent_and_lifecycle():
+    model_args = FakeModelArgs()
+    args = MostBasicAgentArgs(
+        temperature=0.7, use_chain_of_thought=True, chat_model_args=model_args
+    )
+    agent = args.make_agent()
+    assert isinstance(agent, MostBasicAgent)
+    assert agent.temperature == 0.7 and agent.use_chain_of_thought
+    assert agent.chat is model_args.chat
+
+    args.set_reproducibility_mode()
+    assert args.temperature == 0
+
+    assert args.prepare() == "prepared" and model_args.prepared
+    assert args.close() == "closed" and model_args.closed
+
+
+def test_get_action_parses_first_code_block():
+    model_args = FakeModelArgs(responses=["I will click\n```\nclick('a1')\n```\n```\nnoop()\n```"])
+    agent = MostBasicAgentArgs(chat_model_args=model_args).make_agent()
+
+    action, info = agent.get_action(OBS)
+
+    assert action == "click('a1')"
+    assert info.think.startswith("I will click")
+    assert isinstance(info.chat_messages, Discussion)
+    prompt = str(info.chat_messages[1])
+    assert OBS["goal"] in prompt and OBS["pruned_html"] in prompt
+    assert "chain of thoughts" not in prompt
+    assert info.stats["prompt_length"] == 3  # system, user, assistant
+    assert info.stats["response_length"] == len(info.think)
+    assert {"input_tokens", "output_tokens", "cost"} <= set(info.stats)
+    assert info.extra_info["chat_model_args"]["model_name"] == "fake-model"
+
+
+def test_get_action_chain_of_thought_prompt():
+    model_args = FakeModelArgs(responses=["```\nclick('a1')\n```"])
+    agent = MostBasicAgentArgs(use_chain_of_thought=True, chat_model_args=model_args).make_agent()
+    _, info = agent.get_action(OBS)
+    assert "chain of thoughts" in str(info.chat_messages[1])
+
+
+def test_get_action_retries_on_missing_code_block():
+    model_args = FakeModelArgs(responses=["no code here", "```\nclick('a1')\n```"])
+    agent = MostBasicAgentArgs(chat_model_args=model_args).make_agent()
+    action, _ = agent.get_action(OBS)
+    assert action == "click('a1')"
+    assert len(model_args.chat.calls) == 2
+
+
+def test_get_action_raises_after_retries():
+    model_args = FakeModelArgs(responses=["nope"] * 3)
+    agent = MostBasicAgentArgs(chat_model_args=model_args).make_agent()
+    with pytest.raises(ParseError):
+        agent.get_action(OBS)
+
+
+def test_experiment_config():
+    exps = most_basic_agent.experiment_config()
+    assert len(exps) == 2
+    assert [e.agent_args.use_chain_of_thought for e in exps] == [True, False]
+    assert all(e.env_args.task_name == "miniwob.click-button" for e in exps)

tests/agents/test_reproducibility_agent.py

@@ -0,0 +1,130 @@
+from types import SimpleNamespace
+
+import pytest
+from browsergym.experiments.agent import AgentInfo
+from bgym import HighLevelActionSetArgs
+
+from agentlab.agents.generic_agent import AGENT_4o_MINI
+from agentlab.agents.generic_agent.reproducibility_agent import (
+    ReproAgentArgs,
+    ReproChatModel,
+    _diff_stats,
+    _format_messages,
+    _make_agent_stats,
+    _make_backward_compatible,
+    _make_diff,
+    _set_style,
+    make_repro_agent,
+)
+from agentlab.llm.llm_utils import AIMessage, Discussion, HumanMessage, SystemMessage
+
+OLD_MESSAGES = [
+    {"role": "system", "content": "sys"},
+    {"role": "user", "content": "do it"},
+    {"role": "assistant", "content": "<action>click('a1')</action>"},
+]
+
+
+def test_repro_chat_model_replays_next_message():
+    chat = ReproChatModel(OLD_MESSAGES, delay=0)
+    response = chat(OLD_MESSAGES[:2])
+    assert response == OLD_MESSAGES[2]
+    assert chat.new_messages == OLD_MESSAGES
+    assert chat.get_stats() == {}
+
+
+def test_repro_chat_model_does_not_mutate_input():
+    chat = ReproChatModel(OLD_MESSAGES, delay=0)
+    messages = OLD_MESSAGES[:2]
+    chat(messages)
+    assert len(messages) == 2
+
+
+def test_repro_chat_model_exhausted_returns_none_action():
+    chat = ReproChatModel(OLD_MESSAGES[:2], delay=0)
+    response = chat(OLD_MESSAGES[:2])
+    assert response["role"] == "assistant"
+    assert response["content"] == "<action>None</action>"
+
+
+def test_format_messages_list_and_discussion():
+    assert _format_messages(OLD_MESSAGES[:2]) == "system message:\nsys\n\nuser message:\ndo it\n"
+    discussion = Discussion([SystemMessage("sys"), HumanMessage("do it")])
+    assert _format_messages(discussion) == discussion.to_string()
+
+
+def test_diff_stats_identical():
+    assert _diff_stats("a\nb", "a\nb") == dict(lines_added=0, lines_removed=0, difference_ratio=0.0)
+
+
+def test_diff_stats_changes():
+    stats = _diff_stats("a\nb\nc", "a\nB\nc\nd")
+    assert stats["lines_added"] == 2
+    assert stats["lines_removed"] == 1
+    assert stats["difference_ratio"] == pytest.approx(3 / 8)
+
+
+def test_make_diff_has_style_and_no_nowrap():
+    page = _make_diff("same\nold", "same\nnew")
+    assert 'nowrap="nowrap"' not in page
+    assert "<style" in page
+    assert "Old Version" in page and "New Version" in page
+
+
+def test_set_style_appends_and_prepends():
+    html = "<html><head><style>a{}</style></head><body></body></html>"
+    assert "a{}\nb{}" in _set_style(html, "b{}")
+    assert "b{}\na{}" in _set_style(html, "b{}", prepend_previous_style=True)
+
+
+def test_set_style_creates_style_tag():
+    out = _set_style("<html><head></head><body></body></html>", "b{}")
+    assert "<style>\nb{}</style>" in out
+
+
+@pytest.mark.parametrize("agent_info", [{"think": "t"}, AgentInfo(think="t")])
+def test_make_agent_stats(agent_info):
+    old = OLD_MESSAGES
+    new = OLD_MESSAGES[:2] + [AIMessage("<action>noop()</action>")]
+    action, info = _make_agent_stats("noop()", agent_info, None, old, new)
+    assert action == "noop()"
+    assert isinstance(info, AgentInfo)
+    assert info.think == "t"
+    assert "<table" in info.html_page
+    assert info.stats["lines_added"] == 1 and info.stats["lines_removed"] == 1
+
+
+@pytest.mark.parametrize("action_set", ["bid+nav", ["bid", "nav"]])
+def test_make_backward_compatible_converts_action_set(action_set):
+    args = SimpleNamespace(
+        flags=SimpleNamespace(action=SimpleNamespace(action_set=action_set, multi_actions=True))
+    )
+    out = _make_backward_compatible(args)
+    new_set = out.flags.action.action_set
+    assert isinstance(new_set, HighLevelActionSetArgs)
+    assert list(new_set.subsets) == ["bid", "nav"]
+    assert new_set.multiaction is True
+
+
+def test_make_backward_compatible_keeps_new_format():
+    action_set = HighLevelActionSetArgs(subsets=["bid"])
+    args = SimpleNamespace(
+        flags=SimpleNamespace(action=SimpleNamespace(action_set=action_set, multi_actions=False))
+    )
+    assert _make_backward_compatible(args).flags.action.action_set is action_set
+
+
+def test_make_repro_agent(tmp_path):
+    repro = make_repro_agent(AGENT_4o_MINI, tmp_path)
+    assert isinstance(repro, ReproAgentArgs)
+    assert repro._repro_dir == tmp_path
+    assert repro.agent_name.startswith("Repro_")
+    assert repro.chat_model_args is AGENT_4o_MINI.chat_model_args
+    assert repro.flags is AGENT_4o_MINI.flags
+
+
+def test_make_repro_agent_rejects_bad_inputs(tmp_path):
+    with pytest.raises(AssertionError):
+        make_repro_agent(SimpleNamespace(agent_name="x"), tmp_path)
+    with pytest.raises(AssertionError):
+        make_repro_agent(AGENT_4o_MINI, tmp_path / "missing")

tests/analyze/test_archive_studies.py

@@ -0,0 +1,77 @@
+import pandas as pd
+import pytest
+
+from agentlab.analyze.archive_studies import StudyInfo, search_for_reasons_to_archive
+
+
+def _make_study(result_dir, name, rows):
+    study_dir = result_dir / name
+    study_dir.mkdir()
+    pd.DataFrame(rows).to_csv(study_dir / "summary_df.csv", index=False)
+    return study_dir
+
+
+@pytest.fixture
+def result_dir(tmp_path):
+    _make_study(tmp_path, "good_study", [{"n_completed": "10/10", "n_err": 0}])
+    _make_study(tmp_path, "miniwob-tiny-test_study", [{"n_completed": "10/10", "n_err": 0}])
+    _make_study(tmp_path, "error_study", [{"n_completed": "10/10", "n_err": 8}])
+    _make_study(tmp_path, "small_study", [{"n_completed": "2/2", "n_err": 0}])
+    _make_study(tmp_path, "nothing_completed", [{"n_completed": "0/4", "n_err": 0}])
+    _make_study(
+        tmp_path,
+        "multi_agent",
+        [{"n_completed": "4/5", "n_err": 1}, {"n_completed": "5/5", "n_err": 0}],
+    )
+    (tmp_path / "not_a_dir.txt").write_text("ignored")
+    (tmp_path / "no_summary_study").mkdir()
+    return tmp_path
+
+
+def test_search_for_reasons_to_archive(result_dir):
+    infos = search_for_reasons_to_archive(result_dir, min_study_size=5)
+    by_name = {info.study_dir.name: info for info in infos}
+
+    # files and unreadable studies are skipped
+    assert "not_a_dir.txt" not in by_name
+    assert "no_summary_study" not in by_name
+    assert len(by_name) == 6
+    assert all(isinstance(info, StudyInfo) for info in infos)
+
+    assert not by_name["good_study"].should_delete
+    assert by_name["good_study"].reason == ""
+    assert by_name["good_study"].study is None  # no study.pkl.gz in the folder
+
+    assert by_name["miniwob-tiny-test_study"].should_delete
+    assert "Miniwob tiny test" in by_name["miniwob-tiny-test_study"].reason
+
+    assert by_name["error_study"].should_delete
+    assert "Less than 50%" in by_name["error_study"].reason
+
+    assert by_name["small_study"].should_delete
+    assert "Too few tasks" in by_name["small_study"].reason
+
+    assert by_name["nothing_completed"].should_delete
+    assert "No tasks completed" in by_name["nothing_completed"].reason
+
+    # rows are aggregated: 9/10 completed, 1 error -> kept
+    assert not by_name["multi_agent"].should_delete
+
+
+def test_search_uses_most_recent_summary(tmp_path):
+    import time
+
+    study_dir = _make_study(tmp_path, "study", [{"n_completed": "0/10", "n_err": 0}])
+    newer = study_dir / "summary_new.csv"
+    time.sleep(0.05)
+    pd.DataFrame([{"n_completed": "10/10", "n_err": 0}]).to_csv(newer, index=False)
+
+    (info,) = search_for_reasons_to_archive(tmp_path)
+    assert info.summary_df["n_completed"].tolist() == ["10/10"]
+    assert not info.should_delete
+
+
+def test_min_study_size_default_keeps_small_study(tmp_path):
+    _make_study(tmp_path, "small", [{"n_completed": "2/2", "n_err": 0}])
+    (info,) = search_for_reasons_to_archive(tmp_path)
+    assert not info.should_delete