datarobot_genai.drmcp.test_utils.mcp_utils_integration

mcp_utils_integration

integration_test_server_params_with_env

integration_test_server_params_with_env(extra_env: dict[str, str], use_stub: bool = True) -> StdioServerParameters

Return integration test server params with additional environment variables.

Useful for enabling specific tool groups (e.g. {"ENABLE_USE_CASE_TOOLS": "true"}).

Source code in datarobot_genai/drmcp/test_utils/mcp_utils_integration.py

def integration_test_server_params_with_env(
    extra_env: dict[str, str],
    use_stub: bool = True,
) -> StdioServerParameters:
    """Return integration test server params with additional environment variables.

    Useful for enabling specific tool groups (e.g. {"ENABLE_USE_CASE_TOOLS": "true"}).
    """
    params = integration_test_mcp_server_params(use_stub=use_stub)
    env = dict(params.env or {})
    env.update(extra_env)
    return StdioServerParameters(
        command=params.command,
        args=params.args,
        env=env,
    )

aclose_shared_stdio_sessions async

aclose_shared_stdio_sessions() -> None

Close every shared session bound to the current event loop.

Source code in datarobot_genai/drmcp/test_utils/mcp_utils_integration.py

async def aclose_shared_stdio_sessions() -> None:
    """Close every shared session bound to the current event loop."""
    loop = asyncio.get_running_loop()
    for key in [k for k in _shared_sessions if k[0] is loop]:
        await _shared_sessions.pop(key).aclose()

integration_test_mcp_session async

integration_test_mcp_session(server_params: StdioServerParameters | None = None, timeout: int = 60, elicitation_callback: Any | None = None, use_stub: bool = True, shared: bool = True) -> AsyncGenerator[ClientSession, None]

Create and connect a client for the MCP server as a context manager.

By default the underlying stdio server subprocess and client session are shared across tests with the same server configuration (see _SharedStdioSession); exiting the context manager does not close them. Callers that need session-level isolation (a custom elicitation_callback, or tests that mutate server state) get a fresh server: pass shared=False, or any elicitation_callback implies it.

Parameters:

Name	Type	Description	Default
server_params	StdioServerParameters | None	Parameters for configuring the server connection	None
timeout	int	Timeout	60
elicitation_callback	Any | None	Optional callback for handling elicitation requests	None
shared	bool	Reuse one server/session per configuration (default). False forces a fresh, test-scoped server subprocess.	True

Yields

ClientSession: Connected MCP client session

Raises

ConnectionError: If session initialization fails
TimeoutError: If session initialization exceeds timeout

Source code in datarobot_genai/drmcp/test_utils/mcp_utils_integration.py

@contextlib.asynccontextmanager
async def integration_test_mcp_session(
    server_params: StdioServerParameters | None = None,
    timeout: int = 60,
    elicitation_callback: Any | None = None,
    use_stub: bool = True,
    shared: bool = True,
) -> AsyncGenerator[ClientSession, None]:
    """
    Create and connect a client for the MCP server as a context manager.

    By default the underlying stdio server subprocess and client session are shared
    across tests with the same server configuration (see _SharedStdioSession); exiting
    the context manager does not close them. Callers that need session-level isolation
    (a custom elicitation_callback, or tests that mutate server state) get a fresh
    server: pass ``shared=False``, or any elicitation_callback implies it.

    Args:
        server_params: Parameters for configuring the server connection
        timeout: Timeout
        elicitation_callback: Optional callback for handling elicitation requests
        shared: Reuse one server/session per configuration (default). False forces a
            fresh, test-scoped server subprocess.

    Yields
    ------
        ClientSession: Connected MCP client session

    Raises
    ------
        ConnectionError: If session initialization fails
        TimeoutError: If session initialization exceeds timeout
    """
    server_params = server_params or integration_test_mcp_server_params(use_stub=use_stub)

    if shared and elicitation_callback is None:
        holder = await _get_shared_stdio_session(server_params, timeout)
        assert holder.session is not None
        yield holder.session
        return

    try:
        async with stdio_client(server_params) as (read_stream, write_stream):
            async with ClientSession(
                read_stream, write_stream, elicitation_callback=elicitation_callback
            ) as session:
                init_result = await asyncio.wait_for(session.initialize(), timeout=timeout)
                # Store the init result on the session for tests that need to inspect capabilities
                session._init_result = init_result  # type: ignore[attr-defined]
                yield session

    except TimeoutError:
        raise TimeoutError(f"Session initialization timed out after {timeout} seconds")