fix(server.py,-utils.py): fix issue where regions weren't parsed properly (#37)

Ensure that regions are resolved in the following order: 1. from --region param 2. from url 3. from environment variable and throws an error if none of these are available. unit tests were also updated to match.

Author: kyoncal

aws_mcp_proxy/server.py

@@ -47,9 +47,11 @@ async def setup_mcp_mode(local_mcp: FastMCP, args) -> None:
 
     # Validate and determine service
     service = determine_service_name(args.endpoint, args.service)
+    logger.debug('Using service: %s', service)
 
     # Validate and determine region
     region = determine_aws_region(args.endpoint, args.region)
+    logger.debug('Using region: %s', region)
 
     # Get profile
     profile = args.profile
@@ -134,7 +136,7 @@ def parse_args():
     parser.add_argument(
         '--region',
         help='AWS region to use (uses AWS_REGION environment variable if not provided, with final fallback to us-east-1)',
-        default=os.getenv('AWS_REGION', 'us-east-1'),
+        default=None,
     )
 
     parser.add_argument(

aws_mcp_proxy/utils.py

@@ -15,13 +15,18 @@
 """Utility functions for the AWS MCP Proxy."""
 
 import httpx
+import logging
+import os
 import re
 from aws_mcp_proxy.sigv4_helper import create_sigv4_client
 from fastmcp.client.transports import StreamableHttpTransport
 from typing import Dict, Optional
 from urllib.parse import urlparse
 
 
+logger = logging.getLogger(__name__)
+
+
 def create_transport_with_sigv4(
     url: str,
     service: str,
@@ -92,7 +97,7 @@ def determine_service_name(endpoint: str, service: Optional[str] = None) -> str:
     return determined_service
 
 
-def determine_aws_region(endpoint: str, region: Optional[str] = None) -> str:
+def determine_aws_region(endpoint: str, region: Optional[str]) -> str:
     """Validate and determine the AWS region.
 
     Args:
@@ -106,6 +111,7 @@ def determine_aws_region(endpoint: str, region: Optional[str] = None) -> str:
         ValueError: If region cannot be determined
     """
     if region:
+        logger.debug('Region determined through explicit parameter')
         return region
 
     # Parse AWS region from endpoint URL
@@ -114,11 +120,16 @@ def determine_aws_region(endpoint: str, region: Optional[str] = None) -> str:
 
     # Extract region name (pattern: service.region.api.aws or service-name.region.api.aws)
     region_match = re.search(r'\.([a-z0-9-]+)\.api\.aws', hostname)
-    determined_region = region_match.group(1) if region_match else None
-
-    if not determined_region:
-        raise ValueError(
-            f"Could not determine AWS region from endpoint '{endpoint}'. "
-            'Please provide the region explicitly using --region argument.'
-        )
-    return determined_region
+    if region_match:
+        logger.debug('Region determined through endpoint URL')
+        return region_match.group(1)
+
+    environment_region = os.getenv('AWS_REGION')
+    if environment_region:
+        logger.debug('Region determined through environment variable')
+        return environment_region
+
+    raise ValueError(
+        f"Could not determine AWS region from endpoint '{endpoint}' or from environment variable AWS_REGION. "
+        'Please provide the region explicitly using --region argument.'
+    )

tests/unit/test_server.py

@@ -15,9 +15,16 @@
 """Tests for the aws-mcp-proxy Server."""
 
 import pytest
-from aws_mcp_proxy.server import main, parse_args, setup_mcp_mode
+from aws_mcp_proxy.server import (
+    add_retry_middleware,
+    add_tool_filtering_middleware,
+    main,
+    parse_args,
+    setup_mcp_mode,
+)
 from aws_mcp_proxy.sigv4_helper import create_sigv4_client
 from aws_mcp_proxy.utils import determine_service_name
+from fastmcp.server.server import FastMCP
 from unittest.mock import AsyncMock, Mock, patch
 
 
@@ -26,10 +33,22 @@ class TestServer:
 
     @patch('aws_mcp_proxy.server.create_transport_with_sigv4')
     @patch('aws_mcp_proxy.server.FastMCP.as_proxy')
-    async def test_setup_mcp_mode(self, mock_as_proxy, mock_create_transport):
+    @patch('aws_mcp_proxy.server.determine_aws_region')
+    @patch('aws_mcp_proxy.server.determine_service_name')
+    @patch('aws_mcp_proxy.server.add_tool_filtering_middleware')
+    @patch('aws_mcp_proxy.server.add_retry_middleware')
+    async def test_setup_mcp_mode(
+        self,
+        mock_add_retry,
+        mock_add_filtering,
+        mock_determine_service,
+        mock_determine_region,
+        mock_as_proxy,
+        mock_create_transport,
+    ):
         """Test that MCP mode is set up correctly."""
         # Arrange
-        mock_mcp = Mock()
+        local_mcp = Mock(spec=FastMCP)
         mock_args = Mock()
         mock_args.endpoint = 'https://test.example.com'
         mock_args.service = 'test-service'
@@ -38,6 +57,10 @@ async def test_setup_mcp_mode(self, mock_as_proxy, mock_create_transport):
         mock_args.read_only = True
         mock_args.retries = 1
 
+        # Mock return values
+        mock_determine_service.return_value = 'test-service'
+        mock_determine_region.return_value = 'us-east-1'
+
         # Mock the transport and proxy
         mock_transport = Mock()
         mock_create_transport.return_value = mock_transport
@@ -46,25 +69,46 @@ async def test_setup_mcp_mode(self, mock_as_proxy, mock_create_transport):
         mock_as_proxy.return_value = mock_proxy
 
         # Act
-        await setup_mcp_mode(mock_mcp, mock_args)
+        await setup_mcp_mode(local_mcp, mock_args)
 
         # Assert
-        mock_create_transport.assert_called_once()
+        mock_determine_service.assert_called_once_with('https://test.example.com', 'test-service')
+        mock_determine_region.assert_called_once_with('https://test.example.com', 'us-east-1')
+        mock_create_transport.assert_called_once_with(
+            'https://test.example.com', 'test-service', 'us-east-1', None
+        )
         mock_as_proxy.assert_called_once_with(mock_transport)
+        mock_add_filtering.assert_called_once_with(mock_proxy, True)
+        mock_add_retry.assert_called_once_with(mock_proxy, 1)
+        mock_proxy.run_async.assert_called_once()
 
     @patch('aws_mcp_proxy.server.create_transport_with_sigv4')
     @patch('aws_mcp_proxy.server.FastMCP.as_proxy')
-    async def test_setup_mcp_mode_with_tools(self, mock_as_proxy, mock_create_transport):
-        """Test that MCP mode registers tools correctly."""
+    @patch('aws_mcp_proxy.server.determine_aws_region')
+    @patch('aws_mcp_proxy.server.determine_service_name')
+    @patch('aws_mcp_proxy.server.add_tool_filtering_middleware')
+    async def test_setup_mcp_mode_no_retries(
+        self,
+        mock_add_filtering,
+        mock_determine_service,
+        mock_determine_region,
+        mock_as_proxy,
+        mock_create_transport,
+    ):
+        """Test that MCP mode setup without retries doesn't add retry middleware."""
         # Arrange
-        mock_mcp = Mock()
+        local_mcp = Mock(spec=FastMCP)
         mock_args = Mock()
         mock_args.endpoint = 'https://test.example.com'
         mock_args.service = 'test-service'
         mock_args.region = 'us-east-1'
-        mock_args.profile = None
-        mock_args.read_only = True
-        mock_args.retries = 1
+        mock_args.profile = 'test-profile'
+        mock_args.read_only = False
+        mock_args.retries = 0  # No retries
+
+        # Mock return values
+        mock_determine_service.return_value = 'test-service'
+        mock_determine_region.return_value = 'us-east-1'
 
         # Mock the transport and proxy
         mock_transport = Mock()
@@ -74,17 +118,88 @@ async def test_setup_mcp_mode_with_tools(self, mock_as_proxy, mock_create_transp
         mock_as_proxy.return_value = mock_proxy
 
         # Act
-        await setup_mcp_mode(mock_mcp, mock_args)
+        await setup_mcp_mode(local_mcp, mock_args)
 
         # Assert
-        mock_create_transport.assert_called_once()
+        mock_determine_service.assert_called_once_with('https://test.example.com', 'test-service')
+        mock_determine_region.assert_called_once_with('https://test.example.com', 'us-east-1')
+        mock_create_transport.assert_called_once_with(
+            'https://test.example.com', 'test-service', 'us-east-1', 'test-profile'
+        )
         mock_as_proxy.assert_called_once_with(mock_transport)
+        mock_add_filtering.assert_called_once_with(mock_proxy, False)
+        mock_proxy.run_async.assert_called_once()
+
+    def test_add_tool_filtering_middleware(self):
+        """Test that tool filtering middleware is added correctly."""
+        # Arrange
+        mock_mcp = Mock()
+
+        # Act
+        add_tool_filtering_middleware(mock_mcp, read_only=True)
+
+        # Assert
+        mock_mcp.add_middleware.assert_called_once()
+        # Verify that the middleware added is a ToolFilteringMiddleware
+        call_args = mock_mcp.add_middleware.call_args[0][0]
+        from aws_mcp_proxy.middleware.tool_filter import ToolFilteringMiddleware
+
+        assert isinstance(call_args, ToolFilteringMiddleware)
+        assert call_args.read_only is True
+
+    def test_add_retry_middleware(self):
+        """Test that retry middleware is added correctly."""
+        # Arrange
+        mock_mcp = Mock()
+
+        # Act
+        add_retry_middleware(mock_mcp, retries=5)
+
+        # Assert
+        mock_mcp.add_middleware.assert_called_once()
+        # Verify that the middleware added is a RetryMiddleware
+        call_args = mock_mcp.add_middleware.call_args[0][0]
+        from fastmcp.server.middleware.error_handling import RetryMiddleware
+
+        assert isinstance(call_args, RetryMiddleware)
 
     @patch('sys.argv', ['test', 'https://test.example.com'])
     def test_parse_args_default(self):
         """Test parse_args with default arguments."""
         args = parse_args()
         assert args.endpoint == 'https://test.example.com'
+        assert args.service is None
+        assert args.region is None
+        assert args.profile is None
+        assert args.read_only is False
+        assert args.log_level == 'INFO'
+        assert args.retries == 0
+
+    @patch(
+        'sys.argv',
+        [
+            'test',
+            'https://test.example.com',
+            '--service',
+            'custom-service',
+            '--region',
+            'us-west-2',
+            '--read-only',
+            '--log-level',
+            'DEBUG',
+            '--retries',
+            '5',
+        ],
+    )
+    def test_parse_args_custom(self):
+        """Test parse_args with custom arguments."""
+        args = parse_args()
+        assert args.endpoint == 'https://test.example.com'
+        assert args.service == 'custom-service'
+        assert args.region == 'us-west-2'
+        assert args.read_only is True
+        assert args.log_level == 'DEBUG'
+        assert args.retries == 5
 
     @patch('aws_mcp_proxy.server.asyncio.run')
     @patch('sys.argv', ['test', 'https://test.example.com'])

tests/unit/test_utils.py

@@ -153,40 +153,66 @@ def test_validate_service_name_invalid_url_failure(self):
 class TestDetermineRegion:
     """Test cases for determine_aws_region function."""
 
-    def test_determine_region_with_region(self):
+    @patch('os.getenv')
+    def test_determine_region_with_region(self, mock_getenv):
         """Test determination when region is provided."""
         endpoint = 'https://mcp.us-east-1.api.aws/mcp'
         region = 'custom-region'
 
         result = determine_aws_region(endpoint, region)
 
         assert result == region
+        # Environment variable should not be checked when region is provided
+        mock_getenv.assert_not_called()
 
-    def test_determine_region_without_region_success(self):
+    @patch('os.getenv')
+    def test_determine_region_without_region_success(self, mock_getenv):
         """Test determination when region is not provided but can be parsed."""
         endpoint = 'https://mcp.us-east-1.api.aws/mcp'
         expected_region = 'us-east-1'
+        mock_getenv.return_value = None
 
-        result = determine_aws_region(endpoint)
+        result = determine_aws_region(endpoint, None)
 
         assert result == expected_region
+        # Environment variable should not be checked when region can be parsed from endpoint
 
-    def test_determine_region_with_complex_service_name(self):
+    @patch('os.getenv')
+    def test_determine_region_with_complex_service_name(self, mock_getenv):
         """Test parsing region from endpoint with complex service name."""
         endpoint = 'https://eks-mcp-beta.us-west-2.api.aws/mcp'
         expected_region = 'us-west-2'
+        mock_getenv.return_value = None
 
-        result = determine_aws_region(endpoint)
+        result = determine_aws_region(endpoint, None)
 
         assert result == expected_region
+        # Environment variable should not be checked when region can be parsed from endpoint
 
-    def test_determine_region_without_region_failure(self):
+    @patch('os.getenv')
+    def test_determine_region_without_region_failure(self, mock_getenv):
         """Test determination when region cannot be determined."""
         endpoint = 'https://service.example.com'
+        mock_getenv.return_value = None
 
         with pytest.raises(ValueError) as exc_info:
-            determine_aws_region(endpoint)
+            determine_aws_region(endpoint, None)
 
         assert 'Could not determine AWS region' in str(exc_info.value)
         assert endpoint in str(exc_info.value)
         assert '--region argument' in str(exc_info.value)
+        mock_getenv.assert_called_once_with('AWS_REGION')
+
+    @patch('os.getenv')
+    def test_determine_region_from_environment(self, mock_getenv):
+        """Test determination from environment variable when endpoint doesn't contain region."""
+        # Arrange
+        endpoint = 'https://test-service.example.com'
+        mock_getenv.return_value = 'us-west-1'
+
+        # Act
+        result = determine_aws_region(endpoint, None)
+
+        # Assert
+        assert result == 'us-west-1'
+        mock_getenv.assert_called_once_with('AWS_REGION')