Skip to content

Testing

AG2 provides a built-in TestConfig utility in the autogen.beta.testing module to help you write unit tests for your agents. It allows you to mock LLM responses and simulate tool execution scenarios without making actual API calls.

How to mock LLM answers#

To mock LLM answers, you can use TestConfig in place of a standard model configuration. Pass the expected responses as arguments to TestConfig. Each argument represents the mocked response for a sequential turn in the conversation.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
import pytest

from autogen.beta import Agent
from autogen.beta.testing import TestConfig

@pytest.mark.asyncio
async def test_mock_llm_answer():
    # Provide a TestConfig with the mocked string response
    agent = Agent("test_agent")

    # Ask the agent, passing the TestConfig
    res = await agent.ask(
        "Hi!",
        config=TestConfig("This is a mocked response."),
    )

    # The agent returns the mocked response
    assert res.content == "This is a mocked response."

How to test tool execution#

You can also use TestConfig to yield tool calls. This allows you to test both successful tool execution and error handling. By providing a ToolCall as the first response and a string as the final response, you can simulate a complete agent-tool interaction loop.

Success case#

To test a successful tool execution, pass a ToolCall followed by the final answer you expect the LLM to provide after the tool executes.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import pytest

from autogen.beta import Agent
from autogen.beta.events import ToolCall
from autogen.beta.testing import TestConfig

@pytest.mark.asyncio
async def test_tool_success():
    # Define a tool
    def my_tool() -> str:
        return "tool execution result"

    agent = Agent("test_agent", tools=[my_tool])

    # Configure TestConfig to first return a ToolCall, then a final string answer
    test_config = TestConfig(
        ToolCall(name="my_tool"),
        "final result",
    )

    res = await agent.ask("Please use my_tool", config=test_config)

    # After the tool is called and succeeds, the agent returns the second mocked event
    assert res.content == "final result"

Errors#

You can test how your agent reacts when a tool raises an exception, or when an unregistered tool is requested by the LLM.

If a tool raises an exception during execution, it will propagate up to the ask method. You can catch and assert this exception in your tests.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import pytest

from autogen.beta import Agent
from autogen.beta.events import ToolCall
from autogen.beta.testing import TestConfig

@pytest.mark.asyncio
async def test_tool_raise_exc():
    # Define a tool that raises an error
    def failing_tool() -> str:
        raise ValueError("Something went wrong")

    test_config = TestConfig(
        ToolCall(name="failing_tool"),
        "result",
    )

    agent = Agent(
        "test_agent",
        config=test_config,
        tools=[failing_tool],
    )

    with pytest.raises(ValueError, match="Something went wrong"):
        await agent.ask("Hi!")

Tool not found#

If the LLM attempts to call a tool that hasn't been registered with the agent, a ToolNotFoundError is raised.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
import pytest

from autogen.beta import Agent
from autogen.beta.events import ToolCall
from autogen.beta.exceptions import ToolNotFoundError
from autogen.beta.testing import TestConfig

@pytest.mark.asyncio
async def test_tool_not_found():
    # Mock the LLM returning a tool call for "unregistered_tool"
    test_config = TestConfig(ToolCall(name="unregistered_tool"))

    # Agent is created WITHOUT any tools
    agent = Agent("test_agent", config=test_config)

    with pytest.raises(ToolNotFoundError, match="Tool `unregistered_tool` not found"):
        await agent.ask("Hi!")