Increase test coverage across AgentLab

.github/workflows/unit_tests.yml

@@ -66,4 +66,4 @@ jobs:
       - name: Run AgentLab Unit Tests
         env:
           MINIWOB_URL: "http://localhost:8080/miniwob/"
-        run: uv run pytest -n 5 --durations=10 -m 'not pricy' -v tests/
+        run: uv run pytest -n 5 --durations=10 -m 'not pricy' -v --cov=src/agentlab --cov-report=term-missing tests/

pyproject.toml

@@ -65,6 +65,7 @@ dev = [
     "flaky",
     "pytest-xdist",
     "pytest-playwright",
+    "pytest-cov",
 ]
 # tapeagents = [
 #     "tapeagents[converters]",

tests/agents/test_agent_utils.py

@@ -0,0 +1,49 @@
+import numpy as np
+import pytest
+from PIL import Image
+
+from agentlab.agents.agent_utils import draw_click_indicator, draw_mouse_pointer
+
+
+def test_draw_mouse_pointer_returns_image():
+    """Test that draw_mouse_pointer returns a valid RGBA image."""
+    img = Image.new("RGB", (100, 100), color="white")
+    result = draw_mouse_pointer(img, 50, 50)
+    assert isinstance(result, Image.Image)
+    assert result.mode == "RGBA"
+    assert result.size == (100, 100)
+
+
+def test_draw_mouse_pointer_modifies_pixels():
+    """Test that draw_mouse_pointer actually draws something."""
+    img = Image.new("RGB", (100, 100), color="white")
+    result = draw_mouse_pointer(img, 50, 50)
+    arr = np.array(result)
+    # The pointer should have some non-white, non-fully-opaque pixels
+    # At minimum, the alpha channel should have values < 255 somewhere
+    assert not np.all(arr[:, :, 3] == 255) or not np.all(arr[:, :, :3] == 255)
+
+
+def test_draw_mouse_pointer_at_origin():
+    """Test drawing pointer at (0, 0)."""
+    img = Image.new("RGB", (100, 100), color="white")
+    result = draw_mouse_pointer(img, 0, 0)
+    assert result.size == (100, 100)
+
+
+def test_draw_click_indicator_returns_image():
+    """Test that draw_click_indicator returns a valid RGBA image."""
+    img = Image.new("RGB", (100, 100), color="white")
+    result = draw_click_indicator(img, 50, 50)
+    assert isinstance(result, Image.Image)
+    assert result.mode == "RGBA"
+    assert result.size == (100, 100)
+
+
+def test_draw_click_indicator_at_center():
+    """Test that draw_click_indicator draws a crosshair pattern."""
+    img = Image.new("RGB", (200, 200), color="white")
+    result = draw_click_indicator(img, 100, 100)
+    arr = np.array(result)
+    # The crosshair lines should create some non-white pixels
+    assert result.size == (200, 200)

tests/agents/test_debug_agent.py

@@ -0,0 +1,49 @@
+from unittest.mock import patch
+
+import bgym
+from browsergym.experiments.agent import AgentInfo
+
+from agentlab.agents.debug_agent import DebugAgent, DebugAgentArgs
+
+
+def test_debug_agent_args_instantiation():
+    """Test that DebugAgentArgs can be instantiated."""
+    args = DebugAgentArgs()
+    assert args.agent_name == "debug"
+    assert args.action_set_args is not None
+
+
+def test_debug_agent_args_make_agent():
+    """Test that DebugAgentArgs.make_agent() returns a DebugAgent."""
+    args = DebugAgentArgs()
+    agent = args.make_agent()
+    assert isinstance(agent, DebugAgent)
+
+
+def test_debug_agent_get_action():
+    """Test DebugAgent.get_action() with mocked input."""
+    args = DebugAgentArgs()
+    agent = args.make_agent()
+
+    mock_obs = {
+        "axtree_txt": "[1] Click me button",
+        "pruned_html": "<button bid='1'>Click me</button>",
+    }
+
+    with patch("builtins.input", return_value="click('1')"):
+        action, agent_info = agent.get_action(mock_obs)
+
+    assert action == "click('1')"
+    assert isinstance(agent_info, AgentInfo)
+
+
+def test_debug_agent_set_benchmark():
+    """Test that set_benchmark configures correctly for miniwob."""
+    args = DebugAgentArgs()
+
+    class FakeBenchmark:
+        name = "miniwob_tiny_test"
+        high_level_action_set_args = bgym.HighLevelActionSetArgs(subsets=["bid"])
+
+    args.set_benchmark(FakeBenchmark(), demo_mode=False)
+    assert args.use_html is True

tests/agents/test_generic_prompt.py

@@ -1,5 +1,6 @@
 from copy import deepcopy
 
+import numpy as np
 import bgym
 import pytest
 from bgym import HighLevelActionSet, HighLevelActionSetArgs
@@ -194,7 +195,10 @@ def test_shrinking_observation():
 def test_main_prompt_elements_gone_one_at_a_time(flag_name: str, expected_prompts):
 
     if flag_name in ["use_thinking", "obs.use_action_history"]:
-        return  # TODO design new tests for those two flags
+        # These flags interact with history flags (use_think_history, use_action_history).
+        # Disabling use_thinking alone doesn't remove <think> from history when
+        # use_think_history is still True. Dedicated tests below cover these cases.
+        return
 
     # Disable the flag
     flags = deepcopy(ALL_TRUE_FLAGS)
@@ -251,6 +255,143 @@ def test_main_prompt_elements_present():
             assert expected in prompt
 
 
+def _make_screenshot_obs():
+    """Return a copy of base_obs with a mock screenshot (numpy array)."""
+    obs = base_obs.copy()
+    obs["screenshot"] = np.zeros((64, 64, 3), dtype=np.uint8)
+    obs["screenshot_som"] = np.zeros((64, 64, 3), dtype=np.uint8)
+    return obs
+
+
+def test_use_screenshot_adds_image():
+    """When use_screenshot=True the prompt should contain a Screenshot section."""
+    flags = deepcopy(ALL_TRUE_FLAGS)
+    flags.obs.use_screenshot = True
+    flags.obs.use_som = False
+
+    obs_with_screenshot = _make_screenshot_obs()
+    obs_history = [
+        obs_with_screenshot | {"pruned_html": html_template.format(1), "last_action_error": ""},
+        obs_with_screenshot | {"pruned_html": html_template.format(2), "last_action_error": ""},
+        obs_with_screenshot | {"pruned_html": html_template.format(3), "last_action_error": ""},
+    ]
+
+    prompt_maker = MainPrompt(
+        action_set=flags.action.action_set.make_action_set(),
+        obs_history=obs_history,
+        actions=ACTIONS,
+        memories=MEMORIES,
+        thoughts=THOUGHTS,
+        previous_plan="1- think\n2- do it",
+        step=2,
+        flags=flags,
+    )
+    prompt = prompt_maker.prompt
+    prompt_str = prompt.__str__(warn_if_image=False)
+    assert "Screenshot" in prompt_str
+
+
+def test_use_som_adds_annotated_screenshot():
+    """When use_som=True the screenshot section should mention bounding boxes."""
+    flags = deepcopy(ALL_TRUE_FLAGS)
+    flags.obs.use_screenshot = True
+    flags.obs.use_som = True
+
+    obs_with_screenshot = _make_screenshot_obs()
+    obs_history = [
+        obs_with_screenshot | {"pruned_html": html_template.format(1), "last_action_error": ""},
+        obs_with_screenshot | {"pruned_html": html_template.format(2), "last_action_error": ""},
+        obs_with_screenshot | {"pruned_html": html_template.format(3), "last_action_error": ""},
+    ]
+
+    prompt_maker = MainPrompt(
+        action_set=flags.action.action_set.make_action_set(),
+        obs_history=obs_history,
+        actions=ACTIONS,
+        memories=MEMORIES,
+        thoughts=THOUGHTS,
+        previous_plan="1- think\n2- do it",
+        step=2,
+        flags=flags,
+    )
+    prompt = prompt_maker.prompt
+    prompt_str = prompt.__str__(warn_if_image=False)
+    assert "bounding boxes" in prompt_str
+
+
+def test_enable_chat_uses_chat_instructions():
+    """When enable_chat=True the prompt should use ChatInstructions (UI Assistant)."""
+    import time as _time
+
+    flags = deepcopy(ALL_TRUE_FLAGS)
+    flags.enable_chat = True
+
+    # ChatInstructions requires 'timestamp' in chat_messages
+    chat_obs = deepcopy(OBS_HISTORY)
+    for obs in chat_obs:
+        for msg in obs["chat_messages"]:
+            msg["timestamp"] = _time.time()
+
+    prompt_maker = MainPrompt(
+        action_set=flags.action.action_set.make_action_set(),
+        obs_history=chat_obs,
+        actions=ACTIONS,
+        memories=MEMORIES,
+        thoughts=THOUGHTS,
+        previous_plan="1- think\n2- do it",
+        step=2,
+        flags=flags,
+    )
+    prompt_str = str(prompt_maker.prompt)
+    assert "UI Assistant" in prompt_str
+
+
+def test_use_thinking_disabled_removes_think_tags():
+    """When use_thinking=False and use_think_history=False the prompt has no <think> tags."""
+    flags = deepcopy(ALL_TRUE_FLAGS)
+    flags.use_thinking = False
+    flags.obs.use_think_history = False
+
+    prompt_str = str(
+        MainPrompt(
+            action_set=flags.action.action_set.make_action_set(),
+            obs_history=OBS_HISTORY,
+            actions=ACTIONS,
+            memories=MEMORIES,
+            thoughts=THOUGHTS,
+            previous_plan="1- think\n2- do it",
+            step=2,
+            flags=flags,
+        ).prompt
+    )
+    assert "<think>" not in prompt_str
+    assert "</think>" not in prompt_str
+
+
+def test_use_action_history_disabled_removes_action_tags():
+    """When obs.use_action_history=False the history should not contain <action> tags."""
+    flags = deepcopy(ALL_TRUE_FLAGS)
+    flags.obs.use_action_history = False
+    # Also disable the action in concrete/abstract examples which may contain action tags
+    flags.use_concrete_example = False
+    flags.use_abstract_example = False
+
+    prompt_str = str(
+        MainPrompt(
+            action_set=flags.action.action_set.make_action_set(),
+            obs_history=OBS_HISTORY,
+            actions=ACTIONS,
+            memories=MEMORIES,
+            thoughts=THOUGHTS,
+            previous_plan="1- think\n2- do it",
+            step=2,
+            flags=flags,
+        ).prompt
+    )
+    assert "<action>" not in prompt_str
+    assert "click('41')" not in prompt_str
+
+
 if __name__ == "__main__":
     # for debugging
     test_shrinking_observation()

tests/agents/test_most_basic_agent.py

@@ -0,0 +1,82 @@
+import re
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+
+from agentlab.agents.most_basic_agent.most_basic_agent import MostBasicAgent, MostBasicAgentArgs
+from agentlab.analyze import inspect_results
+from agentlab.experiments import launch_exp
+from agentlab.experiments.loop import EnvArgs, ExpArgs
+from agentlab.llm.chat_api import BaseModelArgs
+from agentlab.llm.llm_utils import Discussion
+
+
+class _CheatMiniWoBLLM_HTML:
+    """Cheat LLM that finds bids in pruned HTML format (bid="X" attributes).
+
+    MostBasicAgent passes ``obs["pruned_html"]`` to the LLM and expects the
+    response wrapped in triple-tick code blocks.
+    """
+
+    def __call__(self, messages):
+        if isinstance(messages, Discussion):
+            prompt = messages.to_string()
+        else:
+            prompt = messages[1].get("content", "")
+
+        # pruned_html encodes bids as bid="X" attributes
+        match = re.search(r'bid="(\w+)"[^>]*>.*?button', prompt, re.IGNORECASE | re.DOTALL)
+        if not match:
+            match = re.search(r'bid="(\w+)"', prompt)
+
+        if match:
+            bid = match.group(1)
+            action = f'click("{bid}")'
+        else:
+            raise Exception("Can't find the button's bid in HTML")
+
+        return dict(role="assistant", content=f"I'll click the button.\n```\n{action}\n```\n")
+
+    def get_stats(self):
+        return {}
+
+
+@dataclass
+class _CheatMiniWoBLLMArgs_HTML(BaseModelArgs):
+    model_name: str = "test/cheat_miniwob_html"
+
+    def make_model(self):
+        return _CheatMiniWoBLLM_HTML()
+
+
+def test_most_basic_agent_instantiation():
+    """Test MostBasicAgentArgs can be created and produces a MostBasicAgent."""
+    args = MostBasicAgentArgs(chat_model_args=_CheatMiniWoBLLMArgs_HTML())
+    agent = args.make_agent()
+    assert isinstance(agent, MostBasicAgent)
+
+
+def test_most_basic_agent():
+    """End-to-end test verifying the experiment loop completes without crashing.
+
+    The cheat LLM finds a bid in the pruned HTML and clicks it.  We only verify
+    the experiment ran to completion (no stack_trace), not that the reward is 1.0,
+    because the first bid in the HTML is not guaranteed to be the target button.
+    """
+    exp_args = ExpArgs(
+        agent_args=MostBasicAgentArgs(
+            chat_model_args=_CheatMiniWoBLLMArgs_HTML(),
+        ),
+        env_args=EnvArgs(task_name="miniwob.click-test", task_seed=42, max_steps=2),
+    )
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        launch_exp.run_experiments(
+            1, [exp_args], Path(tmp_dir) / "most_basic_agent_test", parallel_backend="joblib"
+        )
+
+        result_record = inspect_results.load_result_df(tmp_dir, progress_fn=None)
+        assert len(result_record) == 1
+        # The experiment loop should complete without an unhandled exception
+        assert result_record["err_msg"].iloc[0] is None
+        assert result_record["stack_trace"].iloc[0] is None