Back to Repositories

Testing Traffic Control Implementation in OpenHands

This test suite validates the traffic control functionality in the AgentController component of OpenHands, focusing on iteration counting, budget management, and headless mode operation. The tests ensure proper message formatting and error handling for various control scenarios.

Test Coverage Overview

The test suite provides comprehensive coverage of the AgentController’s traffic control mechanisms.

Key areas tested include:

Iteration count validation and message formatting
Budget limit monitoring and decimal precision
Headless mode operation and error messaging
Exception handling and error propagation

Implementation Analysis

The implementation utilizes pytest’s asyncio functionality for testing asynchronous controller methods. The testing approach employs fixture-based setup and mock objects to isolate the traffic control logic.

Notable patterns include:

Async/await testing patterns
Mock injection for error capture
Fixture-based test configuration
Controlled environment simulation

Technical Details

Testing tools and configuration:

pytest-asyncio for async test execution
unittest.mock for component isolation
Custom fixtures for controller setup
InMemoryFileStore for storage simulation
EventStream integration
MagicMock for agent simulation

Best Practices Demonstrated

The test suite exemplifies several testing best practices in Python async testing.

Notable practices include:

Isolated test cases with clear scope
Proper async/await handling
Robust error capture and validation
Comprehensive fixture setup
Clear test case documentation
Precise assertion checking

all-hands-ai/openhands

tests/unit/test_traffic_control.py

            
from unittest.mock import MagicMock

import pytest

from openhands.controller.agent_controller import AgentController
from openhands.core.config import AgentConfig, LLMConfig
from openhands.events import EventStream
from openhands.llm.llm import LLM
from openhands.storage import InMemoryFileStore


@pytest.fixture
def agent_controller():
    llm = LLM(config=LLMConfig())
    agent = MagicMock()
    agent.name = 'test_agent'
    agent.llm = llm
    agent.config = AgentConfig()
    event_stream = EventStream(sid='test', file_store=InMemoryFileStore())
    controller = AgentController(
        agent=agent,
        event_stream=event_stream,
        max_iterations=100,
        max_budget_per_task=10.0,
        sid='test',
        headless_mode=False,
    )
    return controller


@pytest.mark.asyncio
async def test_traffic_control_iteration_message(agent_controller):
    """Test that iteration messages are formatted as integers."""
    # Mock _react_to_exception to capture the error
    error = None

    async def mock_react_to_exception(e):
        nonlocal error
        error = e

    agent_controller._react_to_exception = mock_react_to_exception

    await agent_controller._handle_traffic_control('iteration', 200.0, 100.0)
    assert error is not None
    assert 'Current iteration: 200, max iteration: 100' in str(error)


@pytest.mark.asyncio
async def test_traffic_control_budget_message(agent_controller):
    """Test that budget messages keep decimal points."""
    # Mock _react_to_exception to capture the error
    error = None

    async def mock_react_to_exception(e):
        nonlocal error
        error = e

    agent_controller._react_to_exception = mock_react_to_exception

    await agent_controller._handle_traffic_control('budget', 15.75, 10.0)
    assert error is not None
    assert 'Current budget: 15.75, max budget: 10.00' in str(error)


@pytest.mark.asyncio
async def test_traffic_control_headless_mode(agent_controller):
    """Test that headless mode messages are formatted correctly."""
    # Mock _react_to_exception to capture the error
    error = None

    async def mock_react_to_exception(e):
        nonlocal error
        error = e

    agent_controller._react_to_exception = mock_react_to_exception

    agent_controller.headless_mode = True
    await agent_controller._handle_traffic_control('iteration', 200.0, 100.0)
    assert error is not None
    assert 'in headless mode' in str(error)
    assert 'Current iteration: 200, max iteration: 100' in str(error)