Handle Lambda function pagination (#1) (#1310)

Co-authored-by: Danilo Poccia <[email protected]>
Co-authored-by: Scott Schreckengaust <[email protected]>

Author: 3 people

src/lambda-tool-mcp-server/awslabs/lambda_tool_mcp_server/server.py

@@ -264,14 +264,24 @@ def filter_functions_by_tag(functions, tag_key, tag_value):
     return tagged_functions
 
 
+def get_all_lambda_functions():
+    """Retrieve all available Lambda functions using pagination."""
+    paginator = lambda_client.get_paginator('list_functions')
+    all_functions = []
+
+    for page in paginator.paginate():
+        all_functions.extend(page.get('Functions', []))
+
+    return all_functions
+
+
 def register_lambda_functions():
     """Register Lambda functions as individual tools."""
     try:
         logger.info('Registering Lambda functions as individual tools...')
-        functions = lambda_client.list_functions()
 
         # Get all functions
-        all_functions = functions['Functions']
+        all_functions = get_all_lambda_functions()
         logger.info(f'Total Lambda functions found: {len(all_functions)}')
 
         # First filter by function name if prefix or list is set

src/lambda-tool-mcp-server/tests/conftest.py

@@ -10,31 +10,35 @@ def mock_lambda_client():
     """Create a mock boto3 Lambda client."""
     mock_client = MagicMock()
 
-    # Mock list_functions response
-    mock_client.list_functions.return_value = {
-        'Functions': [
-            {
-                'FunctionName': 'test-function-1',
-                'FunctionArn': 'arn:aws:lambda:us-east-1:123456789012:function:test-function-1',
-                'Description': 'Test function 1 description',
-            },
-            {
-                'FunctionName': 'test-function-2',
-                'FunctionArn': 'arn:aws:lambda:us-east-1:123456789012:function:test-function-2',
-                'Description': 'Test function 2 description',
-            },
-            {
-                'FunctionName': 'prefix-test-function-3',
-                'FunctionArn': 'arn:aws:lambda:us-east-1:123456789012:function:prefix-test-function-3',
-                'Description': 'Test function 3 with prefix',
-            },
-            {
-                'FunctionName': 'other-function',
-                'FunctionArn': 'arn:aws:lambda:us-east-1:123456789012:function:other-function',
-                'Description': '',  # Empty description
-            },
-        ]
-    }
+    # Mock list_functions paginator response
+    paginator_mock = MagicMock()
+    paginator_mock.paginate.return_value = [
+        {
+            'Functions': [
+                {
+                    'FunctionName': 'test-function-1',
+                    'FunctionArn': 'arn:aws:lambda:us-east-1:123456789012:function:test-function-1',
+                    'Description': 'Test function 1 description',
+                },
+                {
+                    'FunctionName': 'test-function-2',
+                    'FunctionArn': 'arn:aws:lambda:us-east-1:123456789012:function:test-function-2',
+                    'Description': 'Test function 2 description',
+                },
+                {
+                    'FunctionName': 'prefix-test-function-3',
+                    'FunctionArn': 'arn:aws:lambda:us-east-1:123456789012:function:prefix-test-function-3',
+                    'Description': 'Test function 3 with prefix',
+                },
+                {
+                    'FunctionName': 'other-function',
+                    'FunctionArn': 'arn:aws:lambda:us-east-1:123456789012:function:other-function',
+                    'Description': '',  # Empty description
+                },
+            ]
+        }
+    ]
+    mock_client.get_paginator.return_value = paginator_mock
 
     # Mock list_tags response
     def mock_list_tags(Resource):

src/lambda-tool-mcp-server/tests/test_integration.py

@@ -34,15 +34,17 @@ def test_mcp_initialization(self, mock_lambda_client):
         def test_tool_registration(self, mock_lambda_client, mock_create_lambda_tool):
             """Test that Lambda functions are registered as tools."""
             # Set up the mock
-            mock_lambda_client.list_functions.return_value = {
-                'Functions': [
-                    {
-                        'FunctionName': 'test-function',
-                        'FunctionArn': 'arn:aws:lambda:us-east-1:123456789012:function:test-function',
-                        'Description': 'Test function description',
-                    },
-                ]
-            }
+            mock_lambda_client.get_paginator.return_value.paginate.return_value = [
+                {
+                    'Functions': [
+                        {
+                            'FunctionName': 'test-function',
+                            'FunctionArn': 'arn:aws:lambda:us-east-1:123456789012:function:test-function',
+                            'Description': 'Test function description',
+                        },
+                    ]
+                }
+            ]
 
             # Call the function
             register_lambda_functions()

src/lambda-tool-mcp-server/tests/test_server.py

@@ -416,8 +416,8 @@ def test_register_with_no_filters(self, mock_create_lambda_tool, mock_lambda_cli
         @patch('awslabs.lambda_tool_mcp_server.server.lambda_client')
         def test_register_error_handling(self, mock_lambda_client):
             """Test error handling in register_lambda_functions."""
-            # Make list_functions raise an exception
-            mock_lambda_client.list_functions.side_effect = Exception('Error listing functions')
+            # Make get_paginator raise an exception
+            mock_lambda_client.get_paginator.side_effect = Exception('Error listing functions')
 
             # Should not raise an exception
             register_lambda_functions()