Skip to content

Index

Public testing helpers for MADA and extension packages.

The objects re-exported here are intended to be stable import targets for test code outside the core repository. They package up the server-state assertions, tool-discovery helpers, and AI-driven end-to-end runner used by the MADA test suite so extension packages can reuse the same utilities instead of copying them into their own tests/ directories.

AgentTestRunner

Async test harness for starting MCP servers and querying them through an agent.

The runner accepts explicit paths to a server configuration and an agent configuration. During startup it randomizes server ports, rewrites matching MCP URLs in the agent config, launches the configured servers, validates the observed server state, and then initializes the requested agent class.

The default agent class is MultiServerAgent, but tests may substitute any compatible implementation that accepts config_path=..., exposes an async initialize(stack) method, and supports process_query(...).

Source code in mada_tools/testing/agent_runner.py 
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
class AgentTestRunner:
    """Async test harness for starting MCP servers and querying them through an agent.

    The runner accepts explicit paths to a server configuration and an agent
    configuration. During startup it randomizes server ports, rewrites matching
    MCP URLs in the agent config, launches the configured servers, validates the
    observed server state, and then initializes the requested agent class.

    The default agent class is `MultiServerAgent`, but tests may substitute any
    compatible implementation that accepts `config_path=...`, exposes an async
    `initialize(stack)` method, and supports `process_query(...)`.
    """

    def __init__(
        self,
        servers_config_path: Path,
        agent_config_path: Path,
        agent_cls: type[AgentProtocol] = MultiServerAgent,
    ):
        """Initialize the test runner.

        Args:
            servers_config_path: Path to the MCP server configuration JSON.
            agent_config_path: Path to the agent configuration JSON.
            agent_cls: Agent implementation to instantiate after servers are
                running. Defaults to `MultiServerAgent`.
        """
        self.base_servers_config_path = Path(servers_config_path)
        self.base_agent_config_path = Path(agent_config_path)
        self.servers_config_path = self.base_servers_config_path
        self.agent_config_path = self.base_agent_config_path
        self.agent_cls = agent_cls

        self.server_manager = ServerManager(
            state_file=Path.home() / ".mada" / f"server_statuses_{uuid.uuid4().hex}.json"
        )
        self.stack: AsyncExitStack | None = None
        self.agent: AgentProtocol | None = None
        self.servers_config: dict[str, Any] | None = None
        self._generated_config_paths: list[Path] = []

    async def __aenter__(self):
        """Start managed resources when entering an async context."""
        await self.start()
        return self

    async def __aexit__(self, exc_type, exc, tb):
        """Always tear down managed resources when leaving an async context."""
        await self.close()

    async def start(self):
        """Prepare randomized configs, start servers, and initialize the agent.

        Raises:
            FileNotFoundError: If either input configuration path does not
                exist.
            Exception: Propagates any server-startup or agent-initialization
                error after best-effort cleanup.
        """
        if not self.base_servers_config_path.exists():
            raise FileNotFoundError(f"Servers config not found: {self.base_servers_config_path}")

        if not self.base_agent_config_path.exists():
            raise FileNotFoundError(f"Agent config not found: {self.base_agent_config_path}")

        servers_data = json.loads(self.base_servers_config_path.read_text(encoding="utf-8"))
        agent_data = json.loads(self.base_agent_config_path.read_text(encoding="utf-8"))

        port_map = {}

        for name, server in servers_data.get("servers", {}).items():
            if "port" in server:
                port_map[name] = random.randint(1024, 65535)
                server["port"] = port_map[name]

        for name, mcp_server in agent_data.get("mcp_servers", {}).items():
            if name in port_map and "url" in mcp_server:
                mcp_server["url"] = f"http://localhost:{port_map[name]}/mcp"

        self.servers_config_path = self.base_servers_config_path.with_name(
            self.base_servers_config_path.stem + f"_randomized_{uuid.uuid4().hex}.json"
        )
        self.agent_config_path = self.base_agent_config_path.with_name(
            self.base_agent_config_path.stem + f"_randomized_{uuid.uuid4().hex}.json"
        )
        self._generated_config_paths = [self.servers_config_path, self.agent_config_path]

        self.servers_config_path.write_text(json.dumps(servers_data, indent=2), encoding="utf-8")
        self.agent_config_path.write_text(json.dumps(agent_data, indent=2), encoding="utf-8")

        self.servers_config = load_server_config(self.servers_config_path)
        self.server_manager.start_servers(self.servers_config_path)

        try:
            active_servers = self.server_manager.state_manager.get_servers(validate=True)
            validate_server_state(self.servers_config["servers"], active_servers)

            self.agent = self.agent_cls(config_path=str(self.agent_config_path))
            self.stack = AsyncExitStack()
            await self.stack.__aenter__()
            await self.agent.initialize(self.stack)

        except Exception:
            await self.close()
            raise

    async def process_query(self, prompt: str) -> str:
        """Process one prompt against the initialized agent.

        Args:
            prompt: Natural-language prompt to send through the managed agent.

        Returns:
            str: Agent response including tool-context annotations.
        """
        assert self.agent is not None, "Call start() first"
        return await self.agent.process_query(prompt, add_tool_context=True)

    async def close(self):
        """Tear down the agent, servers, and temporary config files.

        Cleanup is best-effort so tests do not leak background servers or
        randomized config files even when startup or prompt execution fails.
        """
        try:
            if self.stack is not None:
                await self.stack.aclose()
                self.stack = None
        finally:
            self.server_manager.stop_servers()
            for path in self._generated_config_paths:
                try:
                    path.unlink(missing_ok=True)
                except OSError:
                    pass
            self._generated_config_paths = []
            self.servers_config_path = self.base_servers_config_path
            self.agent_config_path = self.base_agent_config_path

__aenter__() async

Start managed resources when entering an async context.

Source code in mada_tools/testing/agent_runner.py 
80
81
82
83
async def __aenter__(self):
    """Start managed resources when entering an async context."""
    await self.start()
    return self

__aexit__(exc_type, exc, tb) async

Always tear down managed resources when leaving an async context.

Source code in mada_tools/testing/agent_runner.py 
85
86
87
async def __aexit__(self, exc_type, exc, tb):
    """Always tear down managed resources when leaving an async context."""
    await self.close()

__init__(servers_config_path, agent_config_path, agent_cls=MultiServerAgent)

Initialize the test runner.

Parameters:

Name Type Description Default
servers_config_path Path

Path to the MCP server configuration JSON.

 required
agent_config_path Path

Path to the agent configuration JSON.

 required
agent_cls type[AgentProtocol]

Agent implementation to instantiate after servers are running. Defaults to MultiServerAgent.

 MultiServerAgent
Source code in mada_tools/testing/agent_runner.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def __init__(
    self,
    servers_config_path: Path,
    agent_config_path: Path,
    agent_cls: type[AgentProtocol] = MultiServerAgent,
):
    """Initialize the test runner.

    Args:
        servers_config_path: Path to the MCP server configuration JSON.
        agent_config_path: Path to the agent configuration JSON.
        agent_cls: Agent implementation to instantiate after servers are
            running. Defaults to `MultiServerAgent`.
    """
    self.base_servers_config_path = Path(servers_config_path)
    self.base_agent_config_path = Path(agent_config_path)
    self.servers_config_path = self.base_servers_config_path
    self.agent_config_path = self.base_agent_config_path
    self.agent_cls = agent_cls

    self.server_manager = ServerManager(
        state_file=Path.home() / ".mada" / f"server_statuses_{uuid.uuid4().hex}.json"
    )
    self.stack: AsyncExitStack | None = None
    self.agent: AgentProtocol | None = None
    self.servers_config: dict[str, Any] | None = None
    self._generated_config_paths: list[Path] = []

close() async

Tear down the agent, servers, and temporary config files.

Cleanup is best-effort so tests do not leak background servers or randomized config files even when startup or prompt execution fails.

Source code in mada_tools/testing/agent_runner.py 
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
async def close(self):
    """Tear down the agent, servers, and temporary config files.

    Cleanup is best-effort so tests do not leak background servers or
    randomized config files even when startup or prompt execution fails.
    """
    try:
        if self.stack is not None:
            await self.stack.aclose()
            self.stack = None
    finally:
        self.server_manager.stop_servers()
        for path in self._generated_config_paths:
            try:
                path.unlink(missing_ok=True)
            except OSError:
                pass
        self._generated_config_paths = []
        self.servers_config_path = self.base_servers_config_path
        self.agent_config_path = self.base_agent_config_path

process_query(prompt) async

Process one prompt against the initialized agent.

Parameters:

Name Type Description Default
prompt str

Natural-language prompt to send through the managed agent.

 required

Returns:

Name Type Description
str str

Agent response including tool-context annotations.
Source code in mada_tools/testing/agent_runner.py 
145
146
147
148
149
150
151
152
153
154
155
async def process_query(self, prompt: str) -> str:
    """Process one prompt against the initialized agent.

    Args:
        prompt: Natural-language prompt to send through the managed agent.

    Returns:
        str: Agent response including tool-context annotations.
    """
    assert self.agent is not None, "Call start() first"
    return await self.agent.process_query(prompt, add_tool_context=True)

start() async

Prepare randomized configs, start servers, and initialize the agent.

Raises:

Type Description
FileNotFoundError

If either input configuration path does not exist.
Exception

Propagates any server-startup or agent-initialization error after best-effort cleanup.
Source code in mada_tools/testing/agent_runner.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
async def start(self):
    """Prepare randomized configs, start servers, and initialize the agent.

    Raises:
        FileNotFoundError: If either input configuration path does not
            exist.
        Exception: Propagates any server-startup or agent-initialization
            error after best-effort cleanup.
    """
    if not self.base_servers_config_path.exists():
        raise FileNotFoundError(f"Servers config not found: {self.base_servers_config_path}")

    if not self.base_agent_config_path.exists():
        raise FileNotFoundError(f"Agent config not found: {self.base_agent_config_path}")

    servers_data = json.loads(self.base_servers_config_path.read_text(encoding="utf-8"))
    agent_data = json.loads(self.base_agent_config_path.read_text(encoding="utf-8"))

    port_map = {}

    for name, server in servers_data.get("servers", {}).items():
        if "port" in server:
            port_map[name] = random.randint(1024, 65535)
            server["port"] = port_map[name]

    for name, mcp_server in agent_data.get("mcp_servers", {}).items():
        if name in port_map and "url" in mcp_server:
            mcp_server["url"] = f"http://localhost:{port_map[name]}/mcp"

    self.servers_config_path = self.base_servers_config_path.with_name(
        self.base_servers_config_path.stem + f"_randomized_{uuid.uuid4().hex}.json"
    )
    self.agent_config_path = self.base_agent_config_path.with_name(
        self.base_agent_config_path.stem + f"_randomized_{uuid.uuid4().hex}.json"
    )
    self._generated_config_paths = [self.servers_config_path, self.agent_config_path]

    self.servers_config_path.write_text(json.dumps(servers_data, indent=2), encoding="utf-8")
    self.agent_config_path.write_text(json.dumps(agent_data, indent=2), encoding="utf-8")

    self.servers_config = load_server_config(self.servers_config_path)
    self.server_manager.start_servers(self.servers_config_path)

    try:
        active_servers = self.server_manager.state_manager.get_servers(validate=True)
        validate_server_state(self.servers_config["servers"], active_servers)

        self.agent = self.agent_cls(config_path=str(self.agent_config_path))
        self.stack = AsyncExitStack()
        await self.stack.__aenter__()
        await self.agent.initialize(self.stack)

    except Exception:
        await self.close()
        raise

collect_server_tools(active_servers, expected_tools) async

Collect tools exposed by each active server and optionally assert on them.

Parameters:

Name Type Description Default
active_servers dict[str, Any]

Active server objects keyed by server name.

 required
expected_tools dict[str, set[str]]

Optional expected tool-name sets keyed by server name. When a server name is present, the discovered tool set must match exactly.

 required

Returns:

Type Description
dict[str, dict[str, Any]]

dict[str, dict[str, Any]]: Per-server connection information and the set
dict[str, dict[str, Any]]

of discovered tool names.
Source code in mada_tools/testing/server_checks.py 
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
async def collect_server_tools(
    active_servers: dict[str, Any],
    expected_tools: dict[str, set[str]],
) -> dict[str, dict[str, Any]]:
    """Collect tools exposed by each active server and optionally assert on them.

    Args:
        active_servers: Active server objects keyed by server name.
        expected_tools: Optional expected tool-name sets keyed by server name.
            When a server name is present, the discovered tool set must match
            exactly.

    Returns:
        dict[str, dict[str, Any]]: Per-server connection information and the set
        of discovered tool names.
    """
    results: dict[str, dict[str, Any]] = {}

    for name, active in active_servers.items():
        url = f"http://{active.host}:{active.port}/mcp"

        async with AsyncExitStack() as stack:
            transport_cm = streamable_http_client(url)
            read_stream, write_stream, _ = await stack.enter_async_context(transport_cm)

            session = ClientSession(read_stream, write_stream)
            await stack.enter_async_context(session)

            await session.initialize()
            tools_result = await session.list_tools()
            actual_tool_names = {tool.name for tool in tools_result.tools}

        if name in expected_tools:
            assert expected_tools[name] == actual_tool_names, (
                f"{name} tool mismatch. Expected {expected_tools[name]}, got {actual_tool_names}"
            )

        results[name] = {
            "host": active.host,
            "port": active.port,
            "tools": actual_tool_names,
        }

    return results

get_server_env_vars(config_path, server_key)

Return the environment-variable mapping for one configured server.

Parameters:

Name Type Description Default
config_path Path

Path to the server configuration JSON file.

 required
server_key str

Key identifying the server entry inside the servers mapping.

 required

Returns:

Type Description
dict[str, Any]

dict[str, Any]: Environment-variable mapping for the requested server,
dict[str, Any]

or an empty dictionary when the server has no env_vars block.

Raises:

Type Description
KeyError

If server_key does not exist in the configuration.
Source code in mada_tools/testing/server_checks.py 
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def get_server_env_vars(config_path: Path, server_key: str) -> dict[str, Any]:
    """Return the environment-variable mapping for one configured server.

    Args:
        config_path: Path to the server configuration JSON file.
        server_key: Key identifying the server entry inside the `servers`
            mapping.

    Returns:
        dict[str, Any]: Environment-variable mapping for the requested server,
        or an empty dictionary when the server has no `env_vars` block.

    Raises:
        KeyError: If `server_key` does not exist in the configuration.
    """
    data = load_server_config(config_path)
    servers = data.get("servers", {})

    if server_key not in servers:
        raise KeyError(f"Server key '{server_key}' not found in config.")

    return servers[server_key].get("env_vars", {})

load_server_config(config_path)

Load and parse a server configuration JSON file.

Parameters:

Name Type Description Default
config_path Path

Path to a JSON file containing a top-level servers mapping.

 required

Returns:

Type Description
dict[str, Any]

dict[str, Any]: Parsed configuration dictionary.

Raises:

Type Description
FileNotFoundError

If the configuration file does not exist.
Source code in mada_tools/testing/server_checks.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def load_server_config(config_path: Path) -> dict[str, Any]:
    """Load and parse a server configuration JSON file.

    Args:
        config_path: Path to a JSON file containing a top-level `servers`
            mapping.

    Returns:
        dict[str, Any]: Parsed configuration dictionary.

    Raises:
        FileNotFoundError: If the configuration file does not exist.
    """
    if not config_path.exists():
        raise FileNotFoundError(f"Config not found: {config_path}")

    with config_path.open("r", encoding="utf-8") as file:
        return json.load(file)

validate_server_state(expected_servers, active_servers)

Validate that active servers match the expected server configuration.

The checks cover presence, host, port, running status, and configured environment variables. The function raises assertion failures directly so it reads naturally inside tests.

Parameters:

Name Type Description Default
expected_servers dict[str, Any]

Expected server definitions from the config file.

 required
active_servers dict[str, Any]

Actual server objects returned by the server-state manager, keyed by server name.

 required
Source code in mada_tools/testing/server_checks.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def validate_server_state(expected_servers: dict[str, Any], active_servers: dict[str, Any]) -> None:
    """Validate that active servers match the expected server configuration.

    The checks cover presence, host, port, running status, and configured
    environment variables. The function raises assertion failures directly so it
    reads naturally inside tests.

    Args:
        expected_servers: Expected server definitions from the config file.
        active_servers: Actual server objects returned by the server-state
            manager, keyed by server name.
    """
    for name, expected in expected_servers.items():
        assert name in active_servers, f"Missing server: {name}"

        active = active_servers[name]
        assert active.name == name
        assert active.host == expected["host"]
        assert active.port == expected["port"]
        assert active.status == ServerStatus.RUNNING

        for key, value in expected.get("env_vars", {}).items():
            assert active.env_vars.get(key) == value, (
                f"{name} env var mismatch for {key}: expected {value}, got {active.env_vars.get(key)}"
            )