Skip to content

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 
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
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,
    )

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) -> AsyncGenerator[ClientSession, None]

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

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
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 
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
@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,
) -> AsyncGenerator[ClientSession, None]:
    """
    Create and connect a client for the MCP server as a context manager.

    Args:
        server_params: Parameters for configuring the server connection
        timeout: Timeout
        elicitation_callback: Optional callback for handling elicitation requests

    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)

    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")