How to Test AI Agents: Unit Testing Tool Calls
AI agents are only as reliable as the tools they call. When an LLM decides to invoke a function—whether it's querying a database, calling an API, or running a calculation—that tool call becomes the bridge between probabilistic reasoning and deterministic execution. Unit testing these tool calls is one of the highest-leverage practices you can adopt when building agentic systems, because it isolates the part of your system where bugs are easiest to introduce and hardest to detect.
What Is Unit Testing for Tool Calls?
Unit testing tool calls means verifying that each individual tool function behaves correctly in isolation: given a specific set of arguments, it returns the expected result, handles errors gracefully, and produces outputs the LLM can reason about. Rather than testing the entire agent end-to-end (which is slow, expensive, and non-deterministic), you test the building blocks the agent depends on.
A "tool call" in most agent frameworks consists of three components:
- The tool definition — the schema, name, and description the LLM sees.
- The tool implementation — the actual function that runs when the LLM invokes the tool.
- The argument parsing — the conversion of the LLM's JSON arguments into native function parameters.
Each of these is independently testable, and each can fail in subtle ways.
Why It Matters
End-to-end agent tests are tempting because they feel realistic, but they have serious drawbacks. LLM outputs are non-deterministic, so a passing test today may fail tomorrow even if your code is correct. They're also slow and costly, since every test run consumes tokens. Most dangerously, when an end-to-end test fails, you often can't tell whether the problem was a bad prompt, a misformatted tool argument, or a bug in the tool itself.
Unit testing tool calls solves all three problems. Tool functions are deterministic, fast, and free to run. When a unit test fails, you know exactly which tool broke and why. This lets you catch regressions the moment they happen, rather than discovering them through flaky integration tests or, worse, in production.
Setting Up a Testable Tool Structure
The first step is structuring your tools so they're easy to test. The key principle is separation: keep the tool's business logic separate from the LLM interaction layer. Your tool function should accept plain arguments and return plain values, with no dependency on the agent framework.
A Simple Tool Example
Let's define a tool that looks up a user's account balance. We'll write it as a pure function first, then wrap it for the agent.
# tools/balance.py
from dataclasses import dataclass
from typing import Optional
@dataclass
class Account:
account_id: str
balance: float
currency: str
# In-memory store for demonstration
_ACCOUNTS: dict[str, Account] = {
"acc-001": Account("acc-001", 1250.50, "USD"),
"acc-002": Account("acc-002", 0.00, "USD"),
"acc-003": Account("acc-003", 99999.99, "EUR"),
}
class AccountNotFoundError(Exception):
pass
def get_account_balance(account_id: str) -> dict:
"""Return the balance for a given account ID.
Args:
account_id: The unique identifier for the account.
Returns:
A dict with account_id, balance, and currency.
Raises:
AccountNotFoundError: If the account does not exist.
"""
if not account_id or not isinstance(account_id, str):
raise ValueError("account_id must be a non-empty string")
account = _ACCOUNTS.get(account_id)
if account is None:
raise AccountNotFoundError(f"No account found with id: {account_id}")
return {
"account_id": account.account_id,
"balance": account.balance,
"currency": account.currency,
}
Notice that get_account_balance knows nothing about LLMs, JSON schemas, or agent frameworks. It's just a function. That's what makes it trivially testable.
Wrapping the Tool for the Agent
The agent-facing wrapper lives in a separate module. It defines the schema and handles serialization, but delegates all logic to the pure function.
# tools/agent_tools.py
from tools.balance import get_account_balance, AccountNotFoundError
BALANCE_TOOL_SCHEMA = {
"type": "function",
"function": {
"name": "get_account_balance",
"description": "Look up the current balance of a user account.",
"parameters": {
"type": "object",
"properties": {
"account_id": {
"type": "string",
"description": "The unique account identifier, e.g. 'acc-001'.",
}
},
"required": ["account_id"],
},
},
}
def handle_balance_tool_call(arguments: dict) -> dict:
"""Execute the balance tool from parsed LLM arguments."""
try:
result = get_account_balance(arguments["account_id"])
return {"status": "ok", "data": result}
except AccountNotFoundError as e:
return {"status": "error", "error": str(e)}
except ValueError as e:
return {"status": "error", "error": str(e)}
except KeyError:
return {"status": "error", "error": "Missing required argument: account_id"}
Writing the Unit Tests
Now we test both layers: the pure function and the agent wrapper. We'll use pytest for its concise syntax and powerful fixtures.
Testing the Pure Tool Function
# tests/test_balance.py
import pytest
from tools.balance import get_account_balance, AccountNotFoundError
class TestGetAccountBalance:
"""Tests for the core business logic."""
def test_returns_balance_for_valid_account(self):
result = get_account_balance("acc-001")
assert result == {
"account_id": "acc-001",
"balance": 1250.50,
"currency": "USD",
}
def test_returns_zero_balance(self):
result = get_account_balance("acc-002")
assert result["balance"] == 0.00
assert result["currency"] == "USD"
def test_raises_for_unknown_account(self):
with pytest.raises(AccountNotFoundError) as exc_info:
get_account_balance("acc-999")
assert "acc-999" in str(exc_info.value)
def test_raises_for_empty_string(self):
with pytest.raises(ValueError):
get_account_balance("")
def test_raises_for_none(self):
with pytest.raises(ValueError):
get_account_balance(None) # type: ignore
def test_raises_for_non_string(self):
with pytest.raises(ValueError):
get_account_balance(12345) # type: ignore
def test_return_type_is_dict(self):
result = get_account_balance("acc-003")
assert isinstance(result, dict)
assert set(result.keys()) == {"account_id", "balance", "currency"}
def test_balance_is_numeric(self):
result = get_account_balance("acc-001")
assert isinstance(result["balance"], (int, float))
These tests cover the happy path, edge cases, error conditions, and type contracts. Each test has a single assertion focus, making failures easy to diagnose.
Testing the Agent Wrapper
The wrapper tests verify that LLM-produced arguments are handled correctly, including malformed or missing inputs.
# tests/test_agent_tools.py
import pytest
from tools.agent_tools import handle_balance_tool_call
class TestHandleBalanceToolCall:
"""Tests for the agent-facing wrapper."""
def test_valid_arguments_return_ok(self):
result = handle_balance_tool_call({"account_id": "acc-001"})
assert result["status"] == "ok"
assert result["data"]["account_id"] == "acc-001"
def test_unknown_account_returns_error_status(self):
result = handle_balance_tool_call({"account_id": "acc-999"})
assert result["status"] == "error"
assert "acc-999" in result["error"]
def test_missing_account_id_returns_error(self):
result = handle_balance_tool_call({})
assert result["status"] == "error"
assert "account_id" in result["error"]
def test_empty_account_id_returns_error(self):
result = handle_balance_tool_call({"account_id": ""})
assert result["status"] == "error"
def test_extra_arguments_are_ignored(self):
result = handle_balance_tool_call({
"account_id": "acc-001",
"unexpected_field": "ignored"
})
assert result["status"] == "ok"
def test_none_arguments_handled(self):
result = handle_balance_tool_call({"account_id": None})
assert result["status"] == "error"
def test_result_is_json_serializable(self):
"""The LLM needs to receive JSON, so the result must serialize."""
import json
result = handle_balance_tool_call({"account_id": "acc-001"})
# Should not raise
json.dumps(result)
Testing Tools with External Dependencies
Real tools often call external APIs or databases. For unit tests, you must mock these dependencies so tests remain fast and deterministic.
Mocking an HTTP API Call
# tools/weather.py
import requests
from typing import Optional
def get_weather(city: str, api_key: str) -> dict:
"""Fetch current weather for a city."""
if not city:
raise ValueError("city must be a non-empty string")
response = requests.get(
"https://api.weather.example.com/current",
params={"city": city},
headers={"Authorization": f"Bearer {api_key}"},
timeout=10,
)
response.raise_for_status()
data = response.json()
return {
"city": city,
"temperature_c": data["temp_c"],
"condition": data["condition"],
}
# tests/test_weather.py
import pytest
from unittest.mock import patch, MagicMock
from tools.weather import get_weather
class TestGetWeather:
@patch("tools.weather.requests.get")
def test_returns_weather_on_success(self, mock_get):
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"temp_c": 22.5,
"condition": "sunny",
}
mock_response.raise_for_status.return_value = None
mock_get.return_value = mock_response
result = get_weather("Berlin", "test-key")
assert result == {
"city": "Berlin",
"temperature_c": 22.5,
"condition": "sunny",
}
mock_get.assert_called_once_with(
"https://api.weather.example.com/current",
params={"city": "Berlin"},
headers={"Authorization": "Bearer test-key"},
timeout=10,
)
@patch("tools.weather.requests.get")
def test_raises_on_http_error(self, mock_get):
import requests as req
mock_response = MagicMock()
mock_response.raise_for_status.side_effect = req.exceptions.HTTPError(
"404 Not Found"
)
mock_get.return_value = mock_response
with pytest.raises(req.exceptions.HTTPError):
get_weather("Nowhere", "test-key")
@patch("tools.weather.requests.get")
def test_raises_on_timeout(self, mock_get):
import requests as req
mock_get.side_effect = req.exceptions.Timeout("Request timed out")
with pytest.raises(req.exceptions.Timeout):
get_weather("Berlin", "test-key")
def test_raises_for_empty_city(self):
with pytest.raises(ValueError):
get_weather("", "test-key")
@patch("tools.weather.requests.get")
def test_correct_params_sent(self, mock_get):
mock_response = MagicMock()
mock_response.json.return_value = {"temp_c": 0, "condition": "rain"}
mock_response.raise_for_status.return_value = None
mock_get.return_value = mock_response
get_weather("Tokyo", "secret")
called_kwargs = mock_get.call_args.kwargs
assert called_kwargs["params"] == {"city": "Tokyo"}
assert called_kwargs["headers"]["Authorization"] == "Bearer secret"
By mocking requests.get, the tests run in milliseconds without any network access. The mock also lets you simulate failure modes—timeouts, HTTP errors, malformed JSON—that would be difficult or impossible to trigger reliably against a real API.
Testing Argument Validation and Schema Compliance
LLMs sometimes produce arguments that don't match your schema. Your tool should validate inputs defensively, and your tests should confirm that validation works.
# tools/transfer.py
from pydantic import BaseModel, ValidationError, Field
class TransferArgs(BaseModel):
from_account: str = Field(..., min_length=1)
to_account: str = Field(..., min_length=1)
amount: float = Field(..., gt=0)
def transfer_funds(arguments: dict) -> dict:
"""Validate and execute a fund transfer."""
try:
args = TransferArgs(**arguments)
except ValidationError as e:
return {"status": "error", "error": e.errors()}
if args.from_account == args.to_account:
return {"status": "error", "error": "Cannot transfer to the same account"}
# ... actual transfer logic ...
return {
"status": "ok",
"from": args.from_account,
"to": args.to_account,
"amount": args.amount,
}
# tests/test_transfer.py
import pytest
from tools.transfer import transfer_funds
class TestTransferFunds:
def test_valid_transfer(self):
result = transfer_funds({
"from_account": "acc-001",
"to_account": "acc-002",
"amount": 100.00,
})
assert result["status"] == "ok"
assert result["amount"] == 100.00
def test_rejects_negative_amount(self):
result = transfer_funds({
"from_account": "acc-001",
"to_account": "acc-002",
"amount": -50.00,
})
assert result["status"] == "error"
def test_rejects_zero_amount(self):
result = transfer_funds({
"from_account": "acc-001",
"to_account": "acc-002",
"amount": 0,
})
assert result["status"] == "error"
def test_rejects_same_account(self):
result = transfer_funds({
"from_account": "acc-001",
"to_account": "acc-001",
"amount": 10.00,
})
assert result["status"] == "error"
assert "same account" in result["error"]
def test_rejects_missing_field(self):
result = transfer_funds({
"from_account": "acc-001",
"to_account": "acc-002",
})
assert result["status"] == "error"
def test_rejects_empty_account_id(self):
result = transfer_funds({
"from_account": "",
"to_account": "acc-002",
"amount": 10.00,
})
assert result["status"] == "error"
def test_accepts_string_amount(self):
"""LLMs sometimes send numbers as strings."""
result = transfer_funds({
"from_account": "acc-001",
"to_account": "acc-002",
"amount": "50.00",
})
assert result["status"] == "ok"
assert result["amount"] == 50.00
Testing Multi-Step Tool Sequences
Some tools internally call other tools. You can test these compositions by mocking the inner tools while testing the orchestrator.
# tools/account_summary.py
from tools.balance import get_account_balance, AccountNotFoundError
def get_account_summary(account_ids: list[str]) -> dict:
"""Fetch balances for multiple accounts and produce a summary."""
results = []
errors = []
for account_id in account_ids:
try:
balance = get_account_balance(account_id)
results.append(balance)
except (AccountNotFoundError, ValueError) as e:
errors.append({"account_id": account_id, "error": str(e)})
total = sum(r["balance"] for r in results)
return {
"accounts": results,
"errors": errors,
"total_balance": total,
"account_count": len(results),
}
# tests/test_account_summary.py
import pytest
from unittest.mock import patch
from tools.account_summary import get_account_summary
class TestGetAccountSummary:
@patch("tools.account_summary.get_account_balance")
def test_summarizes_multiple_accounts(self, mock_balance):
mock_balance.side_effect = [
{"account_id": "acc-001", "balance": 100.0, "currency": "USD"},
{"account_id": "acc-002", "balance": 200.0, "currency": "USD"},
]
result = get_account_summary(["acc-001", "acc-002"])
assert result["account_count"] == 2
assert result["total_balance"] == 300.0
assert result["errors"] == []
@patch("tools.account_summary.get_account_balance")
def test_continues_on_error(self, mock_balance):
from tools.balance import AccountNotFoundError
mock_balance.side_effect = [
{"account_id": "acc-001", "balance": 100.0, "currency": "USD"},
AccountNotFoundError("not found"),
]
result = get_account_summary(["acc-001", "acc-999"])
assert result["account_count"] == 1
assert result["total_balance"] == 100.0
assert len(result["errors"]) == 1
assert result["errors"][0]["account_id"] == "acc-999"
def test_empty_list_returns_zero_summary(self):
result = get_account_summary([])
assert result["account_count"] == 0
assert result["total_balance"] == 0
assert result["accounts"] == []
Best Practices
Keep Tool Functions Pure
Design every tool so its core logic is a pure or near-pure function: deterministic inputs produce deterministic outputs, and all side effects (network calls, file I/O, database access) are injected or mockable. This is the single most important design decision for testability.
Test the Contract, Not the Implementation
Focus your assertions on what the tool returns and what errors it raises, not on how it computes the result internally. This makes your tests resilient to refactoring. If you swap an internal algorithm, your tests should still pass as long as the observable behavior is unchanged.
Cover the Edge Cases LLMs Actually Produce
LLMs are creative in ways that human callers aren't. They send numbers as strings, omit optional fields, include unexpected fields, pass null where a string was expected, and sometimes hallucinate entirely wrong argument names. Your test suite should explicitly include cases for these scenarios. Reviewing real production logs of tool call arguments is an excellent way to discover edge cases you hadn't considered.
Verify JSON Serializability
Every tool result will be serialized to JSON before being sent back to the LLM. A tool that returns a datetime object, a custom class, or a bytes value will work fine in unit tests but fail in production. Always include a test that calls json.dumps() on the result.
Use Fixtures for Shared State
If multiple tests share setup—database connections, mock configurations, sample data—use pytest fixtures to keep tests DRY and readable.
# tests/conftest.py
import pytest
from tools.balance import Account, _ACCOUNTS
@pytest.fixture
def sample_accounts():
return {
"acc-001": Account("acc-001", 1250.50, "USD"),
"acc-002": Account("acc-002", 0.00, "USD"),
}
@pytest.fixture
def mock_api_key():
return "test-api-key-12345"
Run Tests in CI on Every Commit
Unit tests for tools are fast enough to run on every push. Integrate them into your CI pipeline so regressions are caught before they reach a staging environment. A complete tool test suite for a medium-sized agent should run in under five seconds.
Don't Mock What You Don't Own Without Integration Tests
If you mock an external API in unit tests, you're testing your assumptions about that API, not the API itself. Pair your mocked unit tests with a small number of integration tests that hit the real API (run less frequently, perhaps nightly) to catch contract drift.
Conclusion
Unit testing tool calls is the foundation of a reliable AI agent testing strategy. By isolating each tool's logic as a pure function, mocking external dependencies, and covering the full range of inputs an LLM might produce, you create a fast, deterministic safety net that catches bugs at the source. This approach complements—but does not replace—higher-level integration and evaluation tests. The result is a development workflow where you can refactor tools with confidence, add new tools without fear of regressions, and ship agent updates knowing that the deterministic core of your system is verified. Start by extracting one tool into a pure function and writing a dozen tests for it; the payoff in debugging time alone will justify the effort.