From 9ee281980506066905cf54b9f05a5e3381c2ac4f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 13:19:38 +0000 Subject: [PATCH] feature: add unit tests for llm_utils, prompt_templates, and agent_utils Co-Authored-By: sachet.agarwal --- tests/agents/test_agent_utils.py | 52 +++++++ tests/llm/test_llm_utils_extra.py | 219 +++++++++++++++++++++++++++++ tests/llm/test_prompt_templates.py | 46 ++++++ 3 files changed, 317 insertions(+) create mode 100644 tests/agents/test_agent_utils.py create mode 100644 tests/llm/test_llm_utils_extra.py create mode 100644 tests/llm/test_prompt_templates.py diff --git a/tests/agents/test_agent_utils.py b/tests/agents/test_agent_utils.py new file mode 100644 index 000000000..4f3a5c81a --- /dev/null +++ b/tests/agents/test_agent_utils.py @@ -0,0 +1,52 @@ +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, +) + + +def test_draw_mouse_pointer_returns_new_image(): + image = Image.new("RGB", (100, 100), "white") + result = draw_mouse_pointer(image, 50, 50) + assert result is not image + assert result.mode == "RGBA" + assert result.size == (100, 100) + # some pixels changed where the pointer was drawn + assert result.getpixel((55, 55)) != (255, 255, 255, 255) + + +def test_draw_click_indicator_returns_new_image(): + image = Image.new("RGB", (100, 100), "red") + result = draw_click_indicator(image, 50, 50) + assert result.mode == "RGBA" + assert result.size == (100, 100) + # center pixel unchanged (gap), line pixels changed + assert result.getpixel((50, 40)) != (255, 0, 0, 255) + + +def test_draw_arrowhead(): + image = Image.new("RGB", (100, 100), "white") + from PIL import ImageDraw + + draw = ImageDraw.Draw(image) + draw_arrowhead(draw, (10, 10), (60, 60)) + assert image.getpixel((59, 59)) != (255, 255, 255) + + +def test_zoom_webpage_calls_evaluate(): + page = MagicMock() + result = zoom_webpage(page, 2.0) + page.evaluate.assert_called_once() + assert "200" in page.evaluate.call_args[0][0] + assert result is page + + +def test_zoom_webpage_invalid_factor(): + with pytest.raises(ValueError): + zoom_webpage(MagicMock(), 0) diff --git a/tests/llm/test_llm_utils_extra.py b/tests/llm/test_llm_utils_extra.py new file mode 100644 index 000000000..02aa76efa --- /dev/null +++ b/tests/llm/test_llm_utils_extra.py @@ -0,0 +1,219 @@ +from unittest.mock import MagicMock + +import numpy as np +import pytest +from PIL import Image + +from agentlab.llm import llm_utils +from agentlab.llm.llm_utils import ( + AIMessage, + BaseMessage, + Discussion, + HumanMessage, + ParseError, + SystemMessage, + extract_html_tags, + image_to_jpg_base64_url, + image_to_png_base64_url, + messages_to_dict, + parse_html_tags, + parse_html_tags_raise, + retry_multiple, +) + + +def test_messages_to_dict_from_dicts(): + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + discussion = messages_to_dict(messages) + assert isinstance(discussion, Discussion) + assert len(discussion) == 3 + assert discussion[0]["role"] == "system" + + +def test_messages_to_dict_from_strings(): + discussion = messages_to_dict(["raw text"]) + assert discussion[0]["role"] == "" + + +def test_messages_to_dict_unknown_type(): + with pytest.raises(ValueError): + messages_to_dict([42]) + + +def test_extract_html_tags(): + text = "abcclickdef" + content = extract_html_tags(text, ["think", "action"]) + assert content["think"] == ["abc", "def"] + assert content["action"] == ["click"] + + +def test_extract_html_tags_missing_key(): + content = extract_html_tags("no tags here", ["think"]) + assert content == {} + + +def test_parse_html_tags_valid(): + text = "abcclick" + content, valid, retry_message = parse_html_tags(text, keys=("think", "action")) + assert valid + assert content["think"] == "abc" + assert content["action"] == "click" + assert retry_message == "" + + +def test_parse_html_tags_missing_key(): + content, valid, retry_message = parse_html_tags("abc", keys=("think", "action")) + assert not valid + assert "action" in retry_message + + +def test_parse_html_tags_optional_key(): + content, valid, _ = parse_html_tags( + "abc", keys=("think",), optional_keys=("action",) + ) + assert valid + assert "action" not in content + + +def test_parse_html_tags_multiple_without_merge(): + _, valid, retry_message = parse_html_tags("12", keys=("a",)) + assert not valid + assert "multiple instances" in retry_message + + +def test_parse_html_tags_merge_multiple(): + content, valid, _ = parse_html_tags("12", keys=("a",), merge_multiple=True) + assert valid + assert content["a"] == "1\n2" + + +def test_parse_html_tags_raise(): + with pytest.raises(ParseError): + parse_html_tags_raise("abc", keys=("think", "action")) + content = parse_html_tags_raise("abc", keys=("think",)) + assert content["think"] == "abc" + + +def _parser(answer): + return parse_html_tags_raise(answer, keys=("a",)) + + +def test_retry_multiple_success_first_try(): + chat = MagicMock(return_value={"role": "assistant", "content": "ok"}) + + parsed_answers, tries = retry_multiple(chat, [], n_retry=2, parser=_parser, log=False) + assert parsed_answers[0]["a"] == "ok" + assert tries == 0 + assert chat.call_count == 1 + + +def test_retry_multiple_retries_then_succeeds(): + chat = MagicMock( + side_effect=[ + {"role": "assistant", "content": "bad"}, + {"role": "assistant", "content": "ok"}, + ] + ) + + parsed_answers, tries = retry_multiple(chat, [], n_retry=2, parser=_parser, log=False) + assert parsed_answers[0]["a"] == "ok" + assert tries == 1 + assert chat.call_count == 2 + + +def test_retry_multiple_exhausts_retries(): + chat = MagicMock(return_value={"role": "assistant", "content": "bad"}) + + with pytest.raises(ParseError): + retry_multiple(chat, [], n_retry=1, parser=_parser, log=False) + + +def test_image_to_jpg_base64_url_from_numpy(): + image = np.zeros((4, 4, 3), dtype=np.uint8) + url = image_to_jpg_base64_url(image) + assert url.startswith("data:image/jpeg;base64,") + + +def test_image_to_png_base64_url_from_pil(): + image = Image.new("RGB", (4, 4)) + url = image_to_png_base64_url(image) + assert url.startswith("data:image/png;base64,") + + +def test_image_base64_url_rgba_converted(): + image = Image.new("RGBA", (4, 4)) + url = image_to_jpg_base64_url(image) + assert url.startswith("data:image/jpeg;base64,") + + +def test_base_message_str_and_add(): + msg = BaseMessage("user", "hello") + assert msg["role"] == "user" + assert "hello" in str(msg) + msg.add_text("world") + assert len(msg["content"]) == 2 + msg.add_image(Image.new("RGB", (2, 2))) + types = [c["type"] for c in msg["content"]] + assert "image_url" in types + + +def test_message_subclasses_roles(): + assert SystemMessage("s")["role"] == "system" + assert HumanMessage("h")["role"] == "user" + ai = AIMessage("a") + assert ai["role"] == "assistant" + + +def test_discussion_basics(): + discussion = Discussion(SystemMessage("sys")) + discussion.append(HumanMessage("question")) + assert len(discussion) == 2 + assert discussion[0]["role"] == "system" + assert discussion.last_message["role"] == "user" + assert "question" in discussion.to_string() + openai_format = discussion.to_openai() + assert isinstance(openai_format, list) and len(openai_format) == 2 + + +def test_discussion_add_content_and_merge(): + discussion = Discussion([SystemMessage("a")]) + discussion.add_text("b") + discussion.merge() + assert "a" in str(discussion) + assert "b" in str(discussion) + + +def test_generic_call_api_with_retries_success(): + response = MagicMock() + client = MagicMock(return_value=response) + result = llm_utils.generic_call_api_with_retries( + client, + {"messages": []}, + is_response_valid_fn=lambda r: True, + rate_limit_exceptions=(), + api_error_exceptions=(), + ) + assert result is response + assert client.call_count == 1 + + +def test_generic_call_api_with_retries_invalid_then_valid(): + good = MagicMock() + bad = MagicMock() + client = MagicMock(side_effect=[bad, good]) + result = llm_utils.generic_call_api_with_retries( + client, + {"messages": []}, + is_response_valid_fn=lambda r: r is good, + rate_limit_exceptions=(), + api_error_exceptions=(), + max_retries=3, + initial_retry_delay_seconds=0, + max_retry_delay_seconds=0, + ) + assert result is good + assert client.call_count == 2 diff --git a/tests/llm/test_prompt_templates.py b/tests/llm/test_prompt_templates.py new file mode 100644 index 000000000..e5ad02708 --- /dev/null +++ b/tests/llm/test_prompt_templates.py @@ -0,0 +1,46 @@ +import pytest + +from agentlab.llm.prompt_templates import ( + STARCHAT_PROMPT_TEMPLATE, + PromptTemplate, + get_prompt_template, +) + + +def test_format_message_roles(): + template = PromptTemplate(system="S:{input}\n", human="H:{input}\n", ai="A:{input}\n") + assert template.format_message({"role": "system", "content": "x"}) == "S:x\n" + assert template.format_message({"role": "user", "content": "y"}) == "H:y\n" + assert template.format_message({"role": "assistant", "content": "z"}) == "A:z\n" + + +def test_format_message_unsupported_role(): + template = PromptTemplate(system="{input}", human="{input}", ai="{input}") + with pytest.raises(ValueError): + template.format_message({"role": "tool", "content": "x"}) + + +def test_construct_prompt(): + template = PromptTemplate( + system="S:{input}\n", human="H:{input}\n", ai="A:{input}\n", prompt_end="A:" + ) + prompt = template.construct_prompt( + [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + ] + ) + assert prompt == "S:sys\nH:hi\nA:" + + +def test_construct_prompt_invalid_messages(): + template = PromptTemplate(system="{input}", human="{input}", ai="{input}") + with pytest.raises(ValueError): + template.construct_prompt(["not a dict"]) + + +def test_get_prompt_template(): + assert get_prompt_template("starcoder-15b") is STARCHAT_PROMPT_TEMPLATE + assert get_prompt_template("starchat-beta") is STARCHAT_PROMPT_TEMPLATE + with pytest.raises(NotImplementedError): + get_prompt_template("unknown-model")