Integration test MCP server.
This server works standalone (base tools only) or detects and loads
user modules if they exist in the project structure.
When running under stdio there are no HTTP headers, so request_user_dr_sdk()
and get_datarobot_access_token() would raise. We patch token resolution to fall
back to credentials (from env) so integration tests can use DATAROBOT_API_TOKEN
without injecting headers.
detect_user_modules() -> Any
Detect if user modules exist in the project.
Returns
Tuple of (config_factory, credentials_factory, lifecycle, module_paths) or None
Source code in datarobot_genai/drmcp/test_utils/integration_mcp_server.py
| def detect_user_modules() -> Any:
"""
Detect if user modules exist in the project.
Returns
-------
Tuple of (config_factory, credentials_factory, lifecycle, module_paths) or None
"""
# Try to find app directory
# When run from library: won't find it
# When run from project: will find it
current_dir = Path.cwd()
# Look for app in current directory or parent directories
for search_dir in [current_dir, current_dir.parent, current_dir.parent.parent]:
app_dir = search_dir / "app"
app_core_dir = app_dir / "core"
if app_core_dir.exists():
# Found user directory - load user modules
try:
module_paths = [
(str(app_dir / "tools"), "app.tools"),
(str(app_dir / "prompts"), "app.prompts"),
(str(app_dir / "resources"), "app.resources"),
]
return (
get_user_config,
get_user_credentials,
ServerLifecycle(),
module_paths,
)
except ImportError:
# User modules don't exist or can't be imported
pass
return None
|
Run the integration test MCP server.
Source code in datarobot_genai/drmcp/test_utils/integration_mcp_server.py
| def main() -> None:
"""Run the integration test MCP server."""
if os.environ.get("MCP_USE_CLIENT_STUBS", "true") == "true":
_apply_dr_client_stubs()
apply_lineage_manager_stubs()
elif os.environ.get("MCP_SERVER_NAME") == "integration":
_patch_datarobot_token_for_stdio()
# Try to detect and load user modules
user_components = detect_user_modules()
if user_components:
# User modules found - create server with user extensions
config_factory, credentials_factory, lifecycle, module_paths = user_components
server = create_mcp_server(
config_factory=config_factory,
credentials_factory=credentials_factory,
lifecycle=lifecycle,
additional_module_paths=module_paths,
transport="stdio",
load_native_mcp_tools=True,
)
else:
# No user modules - create server with base tools only
server = create_mcp_server(transport="stdio", load_native_mcp_tools=True)
server.run()
|