AWS Model Context Protocol Server

"""Unit tests for the tools module.""" import asyncio from unittest.mock import AsyncMock, MagicMock, patch import pytest from aws_mcp_server.tools import ( ALLOWED_UNIX_COMMANDS, execute_piped_command, is_pipe_command, split_pipe_command, validate_unix_command, ) def test_allowed_unix_commands(): """Test that ALLOWED_UNIX_COMMANDS contains expected commands.""" # Verify that common Unix utilities are in the allowed list common_commands = ["grep", "xargs", "cat", "ls", "wc", "sort", "uniq", "jq"] for cmd in common_commands: assert cmd in ALLOWED_UNIX_COMMANDS def test_validate_unix_command(): """Test the validate_unix_command function.""" # Test valid commands for cmd in ["grep pattern", "ls -la", "wc -l", "cat file.txt"]: assert validate_unix_command(cmd), f"Command should be valid: {cmd}" # Test invalid commands for cmd in ["invalid_cmd", "sudo ls", ""]: assert not validate_unix_command(cmd), f"Command should be invalid: {cmd}" def test_is_pipe_command(): """Test the is_pipe_command function.""" # Test commands with pipes assert is_pipe_command("aws s3 ls | grep bucket") assert is_pipe_command("aws s3api list-buckets | jq '.Buckets[].Name' | sort") # Test commands without pipes assert not is_pipe_command("aws s3 ls") assert not is_pipe_command("aws ec2 describe-instances") # Test commands with pipes in quotes (should not be detected as pipe commands) assert not is_pipe_command("aws s3 ls 's3://my-bucket/file|other'") assert not is_pipe_command('aws ec2 run-instances --user-data "echo hello | grep world"') # Test commands with escaped quotes - these should not confuse the parser assert is_pipe_command('aws s3 ls --query "Name=\\"value\\"" | grep bucket') assert not is_pipe_command('aws s3 ls "s3://my-bucket/file\\"|other"') def test_split_pipe_command(): """Test the split_pipe_command function.""" # Test simple pipe command cmd = "aws s3 ls | grep bucket" result = split_pipe_command(cmd) assert result == ["aws s3 ls", "grep bucket"] # Test multi-pipe command cmd = "aws s3api list-buckets | jq '.Buckets[].Name' | sort" result = split_pipe_command(cmd) assert result == ["aws s3api list-buckets", "jq '.Buckets[].Name'", "sort"] # Test with quoted pipe symbols (should not split inside quotes) cmd = "aws s3 ls 's3://bucket/file|name' | grep 'pattern|other'" result = split_pipe_command(cmd) assert result == ["aws s3 ls 's3://bucket/file|name'", "grep 'pattern|other'"] # Test with double quotes cmd = 'aws s3 ls "s3://bucket/file|name" | grep "pattern|other"' result = split_pipe_command(cmd) assert result == ['aws s3 ls "s3://bucket/file|name"', 'grep "pattern|other"'] # Test with escaped quotes cmd = 'aws s3 ls --query "Name=\\"value\\"" | grep bucket' result = split_pipe_command(cmd) assert result == ['aws s3 ls --query "Name=\\"value\\""', "grep bucket"] # Test with escaped pipe symbol in quotes cmd = 'aws s3 ls "s3://bucket/file\\"|name" | grep pattern' result = split_pipe_command(cmd) assert result == ['aws s3 ls "s3://bucket/file\\"|name"', "grep pattern"] @pytest.mark.asyncio async def test_execute_piped_command_success(): """Test successful execution of a piped command.""" with patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_subprocess: # Mock the first process in the pipe first_process_mock = AsyncMock() first_process_mock.returncode = 0 first_process_mock.communicate.return_value = (b"S3 output", b"") # Mock the second process in the pipe second_process_mock = AsyncMock() second_process_mock.returncode = 0 second_process_mock.communicate.return_value = (b"Filtered output", b"") # Set up the mock to return different values on subsequent calls mock_subprocess.side_effect = [first_process_mock, second_process_mock] result = await execute_piped_command("aws s3 ls | grep bucket") assert result["status"] == "success" assert result["output"] == "Filtered output" # Verify first command was called with correct args mock_subprocess.assert_any_call("aws", "s3", "ls", stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE) # Verify second command was called with correct args mock_subprocess.assert_any_call("grep", "bucket", stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE) @pytest.mark.asyncio async def test_execute_piped_command_error_first_command(): """Test error handling in execute_piped_command when first command fails.""" with patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_subprocess: # Mock a failed first process process_mock = AsyncMock() process_mock.returncode = 1 process_mock.communicate.return_value = (b"", b"Command failed: aws") mock_subprocess.return_value = process_mock result = await execute_piped_command("aws s3 ls | grep bucket") assert result["status"] == "error" assert "Command failed: aws" in result["output"] @pytest.mark.asyncio async def test_execute_piped_command_error_second_command(): """Test error handling in execute_piped_command when second command fails.""" with patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_subprocess: # Mock the first process in the pipe (success) first_process_mock = AsyncMock() first_process_mock.returncode = 0 first_process_mock.communicate.return_value = (b"S3 output", b"") # Mock the second process in the pipe (failure) second_process_mock = AsyncMock() second_process_mock.returncode = 1 second_process_mock.communicate.return_value = (b"", b"Command not found: xyz") # Set up the mock to return different values on subsequent calls mock_subprocess.side_effect = [first_process_mock, second_process_mock] result = await execute_piped_command("aws s3 ls | xyz") assert result["status"] == "error" assert "Command not found: xyz" in result["output"] @pytest.mark.asyncio async def test_execute_piped_command_timeout(): """Test timeout handling in execute_piped_command.""" with patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_subprocess: # Mock a process that times out process_mock = AsyncMock() # Use a properly awaitable mock that raises TimeoutError communicate_mock = AsyncMock(side_effect=asyncio.TimeoutError()) process_mock.communicate = communicate_mock # Use regular MagicMock since kill() is not an async method process_mock.kill = MagicMock() mock_subprocess.return_value = process_mock result = await execute_piped_command("aws s3 ls | grep bucket", timeout=1) assert result["status"] == "error" assert "Command timed out after 1 seconds" in result["output"] process_mock.kill.assert_called_once() @pytest.mark.asyncio async def test_execute_piped_command_exception(): """Test general exception handling in execute_piped_command.""" with patch("asyncio.create_subprocess_exec", side_effect=Exception("Test exception")): result = await execute_piped_command("aws s3 ls | grep bucket") assert result["status"] == "error" assert "Failed to execute command" in result["output"] assert "Test exception" in result["output"] @pytest.mark.asyncio async def test_execute_piped_command_empty_command(): """Test handling of empty commands.""" result = await execute_piped_command("") assert result["status"] == "error" assert "Empty command" in result["output"] @pytest.mark.asyncio async def test_execute_piped_command_timeout_during_final_wait(): """Test timeout handling during wait for the final command in a pipe.""" # This test directly tests the branch where a timeout occurs during awaiting the final command with patch("asyncio.wait_for", side_effect=asyncio.TimeoutError()): with patch("aws_mcp_server.tools.split_pipe_command") as mock_split: mock_split.return_value = ["aws s3 ls", "grep bucket"] # We don't need to mock the subprocess - it won't reach that point # because wait_for will raise a TimeoutError first result = await execute_piped_command("aws s3 ls | grep bucket", timeout=5) assert result["status"] == "error" assert "Command timed out after 5 seconds" in result["output"] @pytest.mark.asyncio async def test_execute_piped_command_kill_error_during_timeout(): """Test error handling when killing a process after timeout fails.""" with patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_subprocess: # Mock a process that times out process_mock = AsyncMock() process_mock.communicate.side_effect = asyncio.TimeoutError() process_mock.kill = MagicMock(side_effect=Exception("Failed to kill process")) mock_subprocess.return_value = process_mock result = await execute_piped_command("aws s3 ls", timeout=1) assert result["status"] == "error" assert "Command timed out after 1 seconds" in result["output"] process_mock.kill.assert_called_once() @pytest.mark.asyncio async def test_execute_piped_command_large_output(): """Test output truncation in execute_piped_command.""" from aws_mcp_server.config import MAX_OUTPUT_SIZE with patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_subprocess: # Mock a process with large output process_mock = AsyncMock() process_mock.returncode = 0 # Generate output larger than MAX_OUTPUT_SIZE large_output = "x" * (MAX_OUTPUT_SIZE + 1000) process_mock.communicate.return_value = (large_output.encode("utf-8"), b"") mock_subprocess.return_value = process_mock result = await execute_piped_command("aws s3 ls") assert result["status"] == "success" assert len(result["output"]) <= MAX_OUTPUT_SIZE + 100 # Allow for truncation message assert "output truncated" in result["output"]