Create Simple MCP Server for Integ tests (#29)

Author: JacquesPotgieter

.gitignore

@@ -69,3 +69,6 @@ build
 
 # Documentation
 /doc/_apidoc/
+
+# Ignoring AgentCore config files for Simple MCP Server
+/tests/integ/mcp/simple_mcp_server/*

pyproject.toml

@@ -138,6 +138,8 @@ packages = ["aws_mcp_proxy"]
 exclude_dirs = ["venv", ".venv", "tests"]
 
 [tool.pytest.ini_options]
+log_cli = true
+log_cli_level = "INFO"
 python_files = "test_*.py"
 python_classes = "Test*"
 python_functions = "test_*"

tests/integ/README.md

@@ -0,0 +1,67 @@
+## Integration tests
+
+This folder contains the Integration tests for the aws-mcp-proxy. To help with testing, we have also created a simple MCP Server and MCP Client to mimic a real scenario.
+
+* `mcp/simple_mcp_server/mcp_server.py` - A simple MCP Server which supports different features to mimic the Customers remote MCP Server
+* `mcp/simple_mcp_client.py` - A simple MCP Client which uses the proxy to connect to a remote HTTP MCP Server
+* `test_proxy_simple_mcp_server.py` - Actual tests which uses the above to validate the Proxy is working correctly
+
+## How to use?
+
+These tests can be run against two types of Remote MCP Servers
+
+1. Hosted in Bedrock AgentCore Runtime
+1. Against a Remote URL endpoint
+
+#### Hosted in Bedrock AgentCore Runtime
+
+The Simple MCP Server is ready to easily install on AgentCore Bedrock. It is recommended to follow this testing path to ensure sigv4 is working correctly.
+
+The following instructions to install the Simple MCP Server came from [[Deploy MCP servers in AgentCore Runtime](https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/runtime-mcp.html#runtime-mcp-create-server)].
+
+```bash
+# Update Simple MCP Server to use stateless-http
+# - Edit mcp/simple_mcp_server/mcp_server.py
+# - Change "stateless_http=False" to "stateless_http=True"
+
+# Install AgentCore CLI
+pip install bedrock-agentcore-starter-toolkit
+
+# CD into simple_mcp_server folder
+cd integ/mcp/simple_mcp_server
+
+# Create ECR image from simple server
+# The default options will work (OAuth is not required)
+agentcore configure -e mcp_server.py --protocol MCP
+
+# Upload ECR image and host MCP Server
+agentcore launch --auto-update-on-conflict
+
+# Set the ARN which was posted to the Terminal
+export AGENTCORE_RUNTIME_ARN={newly-created-mcp-server}
+```
+
+Run test against the AgentCore hosted MCP Server
+
+```bash
+uv run pytest -m integ
+```
+
+
+#### Against a Remote URL endpoint
+
+To make testing locally faster, you can also run tests against a remote URL. Since this endpoint might not be hosted on AWS, the sigv4 code path might not be fully tested.
+
+Only use this testing path to ensure the MCP Features are working correctly.
+
+```bash
+# Run MCP Server locally
+uv run tests/integ/mcp/simple_mcp_server/mcp_server.py
+
+# Set endpoint to test against
+export REMOTE_ENDPOINT_URL=http://127.0.0.1:8000/mcp
+```
+
+```bash
+uv run pytest -m integ
+```

tests/integ/conftest.py

@@ -0,0 +1,88 @@
+import logging
+import os
+import pytest
+import pytest_asyncio
+from .mcp.simple_mcp_client import build_mcp_client
+from typing import TypedDict
+
+
+logger = logging.getLogger(__name__)
+
+
+class RemoteMCPServerConfiguration(TypedDict):
+    """Remote MCP Server endpoint config."""
+
+    endpoint: str
+    region_name: str
+
+
+@pytest_asyncio.fixture(loop_scope='module', scope='module')
+async def mcp_client(
+    remote_mcp_server_configuration: RemoteMCPServerConfiguration,
+):
+    """Create MCP Client fixture for using in Integ tests."""
+    client = build_mcp_client(
+        endpoint=remote_mcp_server_configuration['endpoint'],
+        region_name=remote_mcp_server_configuration['region_name'],
+    )
+
+    async with client:
+        yield client
+
+
+@pytest.fixture(scope='module')
+async def is_using_agentcore():
+    """Boolean param if we are currently running against AgentCore Runtime."""
+    if os.environ.get('AGENTCORE_RUNTIME_ARN'):
+        return True
+    else:
+        return False
+
+
+@pytest.fixture(scope='module')
+def remote_mcp_server_configuration(is_using_agentcore: bool):
+    """Configuration to connect to remotely hosted MCP Server."""
+    if is_using_agentcore:
+        logger.info('Will use AgentCore MCP Server')
+        return _build_agent_core_remote_configuration()
+    else:
+        logger.info('Will use remote MCP server defined with ENV variables')
+        return _build_endpoint_environment_remote_configuration()
+
+
+def _build_agent_core_remote_configuration():
+    logger.info('Using AgentCore runtime ARN for remote MCP Server')
+
+    runtime_arn = os.environ.get('AGENTCORE_RUNTIME_ARN')
+    if not runtime_arn:
+        raise RuntimeError('AGENTCORE_RUNTIME_ARN env variable not found')
+
+    agent_core_runtime_url_format = 'https://bedrock-agentcore.{region_name}.amazonaws.com/runtimes/{encoded_arn}/invocations?qualifier=DEFAULT'
+    region_name = runtime_arn.split(':')[3]
+    encoded_arn = runtime_arn.replace(':', '%3A').replace('/', '%2F')
+
+    endpoint = agent_core_runtime_url_format.format(
+        region_name=region_name, encoded_arn=encoded_arn
+    )
+
+    return RemoteMCPServerConfiguration(endpoint=endpoint, region_name=region_name)
+
+
+def _build_endpoint_environment_remote_configuration():
+    logger.info('Using Endpoint environment variable for remote MCP Server')
+
+    remote_endpoint_url = os.environ.get('REMOTE_ENDPOINT_URL')
+    if not remote_endpoint_url:
+        raise RuntimeError('REMOTE_ENDPOINT_URL env variable not found')
+
+    region_name = os.environ.get('AWS_REGION')
+    if not region_name:
+        logger.warn('AWS_REGION param not set. Defaulting to us-east-1')
+        region_name = 'us-east-1'
+
+    logger.info(f'Starting server with config - {remote_endpoint_url=} and {region_name=}')
+
+    return RemoteMCPServerConfiguration(
+        endpoint=remote_endpoint_url,
+        region_name=region_name,
+    )

tests/integ/mcp/__init__.py

@@ -0,0 +1,56 @@
+import boto3
+import fastmcp
+import logging
+from fastmcp.client import StdioTransport
+from fastmcp.client.elicitation import ElicitResult
+
+
+logger = logging.getLogger(__name__)
+
+
+def build_mcp_client(endpoint: str, region_name: str) -> fastmcp.Client:
+    """Create a MCP Client using the aws-mcp-proxy against a remote MCP Server."""
+    return fastmcp.Client(
+        StdioTransport(
+            **_build_mcp_config(
+                endpoint=endpoint,
+                region_name=region_name,
+            )
+        ),
+        elicitation_handler=_basic_elicitation_handler,
+        timeout=10.0,  # seconds
+    )
+
+
+async def _basic_elicitation_handler(message: str, response_type: type, params, context):
+    logger.info(f'Server asks: {message} with response_type {response_type}')
+
+    # Usually the Handler would expect an user Input to control flow via Accept, Decline, Cancel
+    # But in this Integ test we only care that an Elicitation request went through the handler
+    # and responded correctly.
+    # As such, we are explicitly hardcoding the response based on the name of the ResponseType object
+
+    if 'Accept' in response_type.__name__:
+        return response_type(value='Elicitation success')
+
+    if 'Decline' in response_type.__name__:
+        return ElicitResult(action='decline')
+
+    raise RuntimeError(f'Unknown Response-type, rather failing - {response_type}')
+
+
+def _build_mcp_config(endpoint: str, region_name: str):
+    credentials = boto3.Session().get_credentials()
+
+    environment_variables = {
+        'AWS_REGION': region_name,
+        'AWS_ACCESS_KEY_ID': credentials.access_key,
+        'AWS_SECRET_ACCESS_KEY': credentials.secret_key,
+        'AWS_SESSION_TOKEN': credentials.token,
+    }
+
+    return {
+        'command': 'aws-mcp-proxy',
+        'args': [endpoint, '--log-level', 'DEBUG'],
+        'env': environment_variables,
+    }

tests/integ/mcp/simple_mcp_client.py

@@ -0,0 +1,76 @@
+import logging
+from dataclasses import dataclass
+from fastmcp import Context, FastMCP
+from typing import Any
+
+
+logger = logging.getLogger(__name__)
+
+mcp = FastMCP[Any](
+    name='Simple MCP Server',
+    instructions=('Simple MCP Server used in Integ Tests'),
+)
+
+##### Generic Tool Testing
+
+
+@mcp.tool
+def greet(name: str):
+    """MCP Tool which is very simple for testing."""
+    return f'Hello {name}'
+
+
+##### Elicitation Testing
+
+
+@dataclass
+class ElicitationWithAccept:
+    """Class type when requesting Elicitation and expecting it to be accepted."""
+
+    value: str
+
+
+@dataclass
+class ElicitationWithDecline:
+    """Class type when requesting Elicitation and expecting it to be declined."""
+
+    value: str
+
+
+@mcp.tool
+async def elicit_for_my_name(elicitation_expected: str, ctx: Context):
+    """MCP Tool which supports elicitation."""
+    response_type = ElicitationWithAccept
+
+    if 'Decline' in elicitation_expected:
+        response_type = ElicitationWithDecline
+
+    result = await ctx.elicit(message='What is your name?', response_type=response_type)
+
+    if result.action == 'accept':
+        return f'Nice to meet you - {result.data.value}'
+    elif result.action == 'decline':
+        return 'Information not provided'
+    else:
+        return 'cancelled'
+
+
+#### Server Setup
+
+
+def main():
+    """Main entrypoint for running this MCP Server."""
+    logger.info('Starting Simple MCP Server')
+
+    mcp.run(
+        transport='http',
+        host='0.0.0.0',
+        port=8000,
+        # By default, this param is set to False to ensure the Elicitation feature is working
+        # When deploying to AgentCore, this flag must be set to True for MCP to work
+        stateless_http=False,
+    )
+
+
+if __name__ == '__main__':
+    main()

tests/integ/mcp/simple_mcp_server/__init__.py

@@ -0,0 +1,2 @@
+fastmcp
+botocore