diff --git a/03-integrations/elasticache/.env.example b/03-integrations/elasticache/.env.example new file mode 100644 index 00000000..6aa9b7f0 --- /dev/null +++ b/03-integrations/elasticache/.env.example @@ -0,0 +1,40 @@ +# AWS Configuration +AWS_PROFILE=default +AWS_REGION=us-west-2 + +# ElastiCache MCP Server Configuration +# No specific environment variables required for ElastiCache MCP Server +# It uses standard AWS credentials and region settings + +# Valkey MCP Server Configuration +VALKEY_HOST=127.0.0.1 +VALKEY_PORT=6379 +VALKEY_USERNAME= +VALKEY_PWD= +VALKEY_USE_SSL=False +VALKEY_CLUSTER_MODE=False + +# Optional Valkey SSL/TLS Configuration +# VALKEY_CA_PATH= +# VALKEY_SSL_KEYFILE= +# VALKEY_SSL_CERTFILE= +# VALKEY_CERT_REQS=required +# VALKEY_CA_CERTS= + +# Memcached MCP Server Configuration +MEMCACHED_HOST=127.0.0.1 +MEMCACHED_PORT=11211 +MEMCACHED_TIMEOUT=1 +MEMCACHED_CONNECT_TIMEOUT=5 +MEMCACHED_RETRY_TIMEOUT=1 +MEMCACHED_MAX_RETRIES=3 + +# Optional Memcached SSL/TLS Configuration +# MEMCACHED_USE_TLS=true +# MEMCACHED_TLS_CERT_PATH=/path/to/client-cert.pem +# MEMCACHED_TLS_KEY_PATH=/path/to/client-key.pem +# MEMCACHED_TLS_CA_CERT_PATH=/path/to/ca-cert.pem +# MEMCACHED_TLS_VERIFY=true + +# AWS Documentation MCP Server Configuration +AWS_DOCUMENTATION_PARTITION=aws # Use 'aws-cn' for AWS China documentation diff --git a/03-integrations/elasticache/README.md b/03-integrations/elasticache/README.md new file mode 100644 index 00000000..707733bb --- /dev/null +++ b/03-integrations/elasticache/README.md @@ -0,0 +1,189 @@ +# ElastiCache Integration for Strands + +This integration allows you to interact with AWS ElastiCache services using the Strands Agent framework. It supports four different MCP servers: + +1. **ElastiCache MCP Server** - For managing ElastiCache control plane operations (creating/modifying/deleting clusters, etc.) +2. **Valkey MCP Server** - For interacting with Redis/Valkey data in ElastiCache +3. **Memcached MCP Server** - For interacting with Memcached data in ElastiCache +4. **AWS Documentation MCP Server** - For accessing AWS documentation related to ElastiCache and other AWS services + +## Architecture + +The integration uses the Model Context Protocol (MCP) to connect to the different ElastiCache services: + +```mermaid +graph TD + A[Strands Agent] --> B[ElastiCache MCP Server] + A --> C[Valkey MCP Server] + A --> D[Memcached MCP Server] + A --> H[AWS Documentation MCP Server] + B --> E[ElastiCache Control Plane] + C --> F[Valkey Data Plane] + D --> G[Memcached Data Plane] + H --> I[AWS Documentation] + + classDef agent fill:#f9f,stroke:#333,stroke-width:2px; + classDef server fill:#bbf,stroke:#333,stroke-width:1px; + classDef plane fill:#bfb,stroke:#333,stroke-width:1px; + + class A agent; + class B,C,D,H server; + class E,F,G,I plane; +``` + +The diagram above shows how the integration works: + +- **ElastiCache Control Plane**: Manage clusters, replication groups, and serverless caches +- **Valkey Data Plane**: Interact with Redis/Valkey data types (strings, lists, sets, hashes, etc.) +- **Memcached Data Plane**: Interact with Memcached data (key-value operations) +- **AWS Documentation**: Access AWS documentation for ElastiCache and other AWS services + +## Prerequisites + +1. Install `uv` from [Astral](https://docs.astral.sh/uv/getting-started/installation/) or the [GitHub README](https://github.com/astral-sh/uv#installation) +2. Install Python using `uv python install 3.10` +3. Set up AWS credentials with access to AWS ElastiCache services +4. Install the required dependencies: + +```bash +uv venv +source .venv/bin/activate +uv pip install -r requirements.txt +``` + +## Configuration + +Copy the `.env.example` file to `.env` and update the values: + +```bash +cp .env.example .env +``` + +### Environment Variables + +#### AWS Configuration +- `AWS_PROFILE`: AWS credential profile to use +- `AWS_REGION`: AWS region to connect to + +#### Valkey Configuration +- `VALKEY_HOST`: ElastiCache Primary Endpoint or Valkey IP/hostname +- `VALKEY_PORT`: Valkey port (default: 6379) +- `VALKEY_USERNAME`: Database username (optional) +- `VALKEY_PWD`: Database password (optional) +- `VALKEY_USE_SSL`: Enable SSL/TLS (True/False) +- `VALKEY_CLUSTER_MODE`: Enable Valkey Cluster mode (True/False) + +#### Memcached Configuration +- `MEMCACHED_HOST`: Memcached server hostname +- `MEMCACHED_PORT`: Memcached server port (default: 11211) +- `MEMCACHED_TIMEOUT`: Operation timeout in seconds +- `MEMCACHED_CONNECT_TIMEOUT`: Connection timeout in seconds +- `MEMCACHED_RETRY_TIMEOUT`: Retry delay in seconds +- `MEMCACHED_MAX_RETRIES`: Maximum number of retry attempts + +#### AWS Documentation Configuration +- `AWS_DOCUMENTATION_PARTITION`: AWS partition to use for documentation (default: aws, use aws-cn for AWS China) + +#### Strands Model Configuration +- `STRANDS_MODEL_ID`: The ID of the inference profile to use for the Strands agent. Note that you must provide the ID of an inference profile and not the model for this parameter. + +To find the correct inference profile ID, run the following AWS CLI command with the appropriate credentials: + +```bash +aws bedrock list-inference-profiles +``` + +From the output, provide either the `inferenceProfileId` or `inferenceProfileArn` value for the model you want to use. + +## Usage + +Run the integration with a prompt: + +```bash +python main.py "Your prompt here" +``` + +### Options + +- `--mcp-server`: Specify which MCP server to use (elasticache, valkey, memcached, aws-docs, or all) +- `--ro`: Run in read-only mode (no writes allowed) + +### Examples + +1. Use all engines: +```bash +python main.py "Create a new ElastiCache replication group and store some data in it" +``` + +2. Use only the ElastiCache control plane: +```bash +python main.py "List all my ElastiCache clusters" --mcp-server elasticache +``` + +3. Use only Valkey in read-only mode: +```bash +python main.py "Get all keys in my Redis database" --mcp-server valkey --ro +``` + +4. Use only Memcached: +```bash +python main.py "Store a session token in Memcached" --mcp-server memcached +``` + +5. Use only AWS Documentation: +```bash +python main.py "Find documentation about ElastiCache serverless" --mcp-server aws-docs +``` + +## Available Operations + +### ElastiCache Control Plane Operations + +- Create, modify, and delete replication groups +- Create, modify, and delete cache clusters +- Create, modify, and delete serverless caches +- Describe replication groups, cache clusters, and serverless caches +- Manage service updates +- Configure CloudWatch metrics and logs +- Manage cost and usage data + +### AWS Documentation Operations + +- Read AWS documentation pages and convert to markdown +- Search AWS documentation using the official search API +- Get content recommendations for AWS documentation pages +- Get a list of available AWS services (China region only) + +### Valkey Data Operations + +- Strings: SET, GET, APPEND, INCREMENT, etc. +- Lists: LPUSH, RPUSH, LPOP, RPOP, etc. +- Sets: SADD, SMEMBERS, SISMEMBER, etc. +- Sorted Sets: ZADD, ZRANGE, ZRANK, etc. +- Hashes: HSET, HGET, HGETALL, etc. +- Streams: XADD, XREAD, XRANGE, etc. +- JSON: JSON.SET, JSON.GET, JSON.ARRAPPEND, etc. +- Bitmaps: SETBIT, GETBIT, BITCOUNT, etc. +- HyperLogLog: PFADD, PFCOUNT, etc. + +### Memcached Operations + +- GET, SET, ADD, REPLACE +- DELETE, TOUCH, FLUSH +- INCREMENT, DECREMENT +- APPEND, PREPEND +- CAS (Check and Set) +- Multi-operations: GET_MULTI, SET_MULTI, etc. +- Stats and server information + +## Testing + +Run the tests: + +```bash +pytest tests/ +``` + +## License + +This project is licensed under the Apache License 2.0 - see the LICENSE file for details. diff --git a/03-integrations/elasticache/architecture.mmd b/03-integrations/elasticache/architecture.mmd new file mode 100644 index 00000000..8aca945e --- /dev/null +++ b/03-integrations/elasticache/architecture.mmd @@ -0,0 +1,17 @@ +graph TD + A[Strands Agent] --> B[ElastiCache MCP Server] + A --> C[Valkey MCP Server] + A --> D[Memcached MCP Server] + A --> H[AWS Documentation MCP Server] + B --> E[ElastiCache Control Plane] + C --> F[Valkey Data Plane] + D --> G[Memcached Data Plane] + H --> I[AWS Documentation] + + classDef agent fill:#f9f,stroke:#333,stroke-width:2px; + classDef server fill:#bbf,stroke:#333,stroke-width:1px; + classDef plane fill:#bfb,stroke:#333,stroke-width:1px; + + class A agent; + class B,C,D,H server; + class E,F,G,I plane; diff --git a/03-integrations/elasticache/architecture.txt b/03-integrations/elasticache/architecture.txt new file mode 100644 index 00000000..a0e50536 --- /dev/null +++ b/03-integrations/elasticache/architecture.txt @@ -0,0 +1,23 @@ ++-------------------+ + | | + | Strands Agent | + | | + +--------+----------+ + | + | + +--------------------+--------------------+--------------------+ + | | | | + | | | | + +-----------v-----------+ +------v------+ +----------v---------+ +--------v-----------+ + | | | | | | | | + | ElastiCache MCP Server| |Valkey Server| | Memcached Server | | AWS Documentation | + | | | | | | | MCP Server | + +-----------+-----------+ +------+------+ +----------+---------+ +--------+-----------+ + | | | | + | | | | + +-----------v-----------+ +------v------+ +----------v---------+ +--------v-----------+ + | | | | | | | | + | ElastiCache Control | |Valkey Data | | Memcached Data | | AWS Documentation | + | Plane | |Plane | | Plane | | | + | | | | | | | | + +-----------------------+ +-------------+ +--------------------+ +--------------------+ diff --git a/03-integrations/elasticache/main.py b/03-integrations/elasticache/main.py new file mode 100644 index 00000000..04d36d7a --- /dev/null +++ b/03-integrations/elasticache/main.py @@ -0,0 +1,253 @@ +#!/usr/bin/env python3 +import argparse +import os +from strands import Agent +from strands.tools.mcp import MCPClient +from strands_tools import file_read, file_write +from mcp import StdioServerParameters, stdio_client +from dotenv import load_dotenv + +load_dotenv() + + +def get_elasticache_mcp_client(readonly=False): + """Create and return an ElastiCache MCP client.""" + args = ["awslabs.elasticache-mcp-server@latest"] + + # Add readonly flag if specified + if readonly: + args.append("--readonly") + + # Create environment variables dictionary + env = {} + + # AWS Configuration + env["AWS_REGION"] = os.getenv("AWS_REGION", "us-west-2") + if os.getenv("AWS_PROFILE"): + env["AWS_PROFILE"] = os.getenv("AWS_PROFILE") + + # ElastiCache Connection Settings + if os.getenv("ELASTICACHE_MAX_RETRIES"): + env["ELASTICACHE_MAX_RETRIES"] = os.getenv("ELASTICACHE_MAX_RETRIES") + if os.getenv("ELASTICACHE_RETRY_MODE"): + env["ELASTICACHE_RETRY_MODE"] = os.getenv("ELASTICACHE_RETRY_MODE") + if os.getenv("ELASTICACHE_CONNECT_TIMEOUT"): + env["ELASTICACHE_CONNECT_TIMEOUT"] = os.getenv("ELASTICACHE_CONNECT_TIMEOUT") + if os.getenv("ELASTICACHE_READ_TIMEOUT"): + env["ELASTICACHE_READ_TIMEOUT"] = os.getenv("ELASTICACHE_READ_TIMEOUT") + + # Cost Explorer Settings + if os.getenv("COST_EXPLORER_MAX_RETRIES"): + env["COST_EXPLORER_MAX_RETRIES"] = os.getenv("COST_EXPLORER_MAX_RETRIES") + if os.getenv("COST_EXPLORER_RETRY_MODE"): + env["COST_EXPLORER_RETRY_MODE"] = os.getenv("COST_EXPLORER_RETRY_MODE") + if os.getenv("COST_EXPLORER_CONNECT_TIMEOUT"): + env["COST_EXPLORER_CONNECT_TIMEOUT"] = os.getenv("COST_EXPLORER_CONNECT_TIMEOUT") + if os.getenv("COST_EXPLORER_READ_TIMEOUT"): + env["COST_EXPLORER_READ_TIMEOUT"] = os.getenv("COST_EXPLORER_READ_TIMEOUT") + + # CloudWatch Settings + if os.getenv("CLOUDWATCH_MAX_RETRIES"): + env["CLOUDWATCH_MAX_RETRIES"] = os.getenv("CLOUDWATCH_MAX_RETRIES") + if os.getenv("CLOUDWATCH_RETRY_MODE"): + env["CLOUDWATCH_RETRY_MODE"] = os.getenv("CLOUDWATCH_RETRY_MODE") + if os.getenv("CLOUDWATCH_CONNECT_TIMEOUT"): + env["CLOUDWATCH_CONNECT_TIMEOUT"] = os.getenv("CLOUDWATCH_CONNECT_TIMEOUT") + if os.getenv("CLOUDWATCH_READ_TIMEOUT"): + env["CLOUDWATCH_READ_TIMEOUT"] = os.getenv("CLOUDWATCH_READ_TIMEOUT") + + # CloudWatch Logs Settings + if os.getenv("CLOUDWATCH_LOGS_MAX_RETRIES"): + env["CLOUDWATCH_LOGS_MAX_RETRIES"] = os.getenv("CLOUDWATCH_LOGS_MAX_RETRIES") + if os.getenv("CLOUDWATCH_LOGS_RETRY_MODE"): + env["CLOUDWATCH_LOGS_RETRY_MODE"] = os.getenv("CLOUDWATCH_LOGS_RETRY_MODE") + if os.getenv("CLOUDWATCH_LOGS_CONNECT_TIMEOUT"): + env["CLOUDWATCH_LOGS_CONNECT_TIMEOUT"] = os.getenv("CLOUDWATCH_LOGS_CONNECT_TIMEOUT") + if os.getenv("CLOUDWATCH_LOGS_READ_TIMEOUT"): + env["CLOUDWATCH_LOGS_READ_TIMEOUT"] = os.getenv("CLOUDWATCH_LOGS_READ_TIMEOUT") + + # Firehose Settings + if os.getenv("FIREHOSE_MAX_RETRIES"): + env["FIREHOSE_MAX_RETRIES"] = os.getenv("FIREHOSE_MAX_RETRIES") + if os.getenv("FIREHOSE_RETRY_MODE"): + env["FIREHOSE_RETRY_MODE"] = os.getenv("FIREHOSE_RETRY_MODE") + if os.getenv("FIREHOSE_CONNECT_TIMEOUT"): + env["FIREHOSE_CONNECT_TIMEOUT"] = os.getenv("FIREHOSE_CONNECT_TIMEOUT") + if os.getenv("FIREHOSE_READ_TIMEOUT"): + env["FIREHOSE_READ_TIMEOUT"] = os.getenv("FIREHOSE_READ_TIMEOUT") + + return MCPClient( + lambda: stdio_client( + StdioServerParameters( + command="uvx", args=args, env=env + ) + ) + ) + + +def get_valkey_mcp_client(readonly=False): + """Create and return a Valkey MCP client.""" + args = ["awslabs.valkey-mcp-server@latest"] + + # Add readonly flag if specified + if readonly: + args.append("--readonly") + + # Create environment variables dictionary + env = {} + + # Get connection details from environment variables + host = os.getenv("VALKEY_HOST", "127.0.0.1") + port = os.getenv("VALKEY_PORT", "6379") + username = os.getenv("VALKEY_USERNAME") + password = os.getenv("VALKEY_PWD") + use_ssl = os.getenv("VALKEY_USE_SSL", "False").lower() == "true" + cluster_mode = os.getenv("VALKEY_CLUSTER_MODE", "False").lower() == "true" + + # Add environment variables to env dictionary + env["VALKEY_HOST"] = host + env["VALKEY_PORT"] = port + if username: + env["VALKEY_USERNAME"] = username + if password: + env["VALKEY_PWD"] = password + env["VALKEY_USE_SSL"] = "true" if use_ssl else "false" + env["VALKEY_CLUSTER_MODE"] = "true" if cluster_mode else "false" + + return MCPClient( + lambda: stdio_client( + StdioServerParameters( + command="uvx", args=args, env=env + ) + ) + ) + + +def get_memcached_mcp_client(readonly=False): + """Create and return a Memcached MCP client.""" + args = ["awslabs.memcached-mcp-server@latest"] + + # Add readonly flag if specified + if readonly: + args.append("--readonly") + + # Create environment variables dictionary + env = {} + + # Get connection details from environment variables + host = os.getenv("MEMCACHED_HOST", "127.0.0.1") + port = os.getenv("MEMCACHED_PORT", "11211") + + # Add environment variables to env dictionary + env["MEMCACHED_HOST"] = host + env["MEMCACHED_PORT"] = port + + return MCPClient( + lambda: stdio_client( + StdioServerParameters( + command="uvx", args=args, env=env + ) + ) + ) + + +def get_aws_documentation_mcp_client(): + """Create and return an AWS Documentation MCP client.""" + args = ["awslabs.aws-documentation-mcp-server@latest"] + + # Create environment variables dictionary + env = {} + + # Get partition from environment variable or use default + partition = os.getenv("AWS_DOCUMENTATION_PARTITION", "aws") + env["AWS_DOCUMENTATION_PARTITION"] = partition + + return MCPClient( + lambda: stdio_client( + StdioServerParameters( + command="uvx", args=args, env=env + ) + ) + ) + + +def main(): + # Set up argument parser + parser = argparse.ArgumentParser(description="ElastiCache client using strands Agent") + parser.add_argument("prompt", help="Prompt to send to the agent") + parser.add_argument( + "--mcp-server", + choices=["elasticache", "valkey", "memcached", "aws-docs", "all"], + default="all", + help="Specify which MCP server to use (elasticache, valkey, memcached, aws-docs, or all)" + ) + parser.add_argument( + "--ro", action="store_true", help="Run in read-only mode (no writes allowed)" + ) + + parser.add_argument( + "--model", + default=os.getenv("STRANDS_MODEL_ID", "us.anthropic.claude-3-7-sonnet-20250219-v1:0"), + help="Model ID to use for the agent (default: from STRANDS_MODEL_ID env var or Claude 3 Sonnet)" + ) + + # Parse arguments + args = parser.parse_args() + + # Initialize clients based on the selected MCP server + clients = [] + + if args.mcp_server in ["elasticache", "all"]: + try: + elasticache_client = get_elasticache_mcp_client(readonly=args.ro) + clients.append(elasticache_client) + except Exception as e: + print(f"Warning: Failed to initialize ElastiCache client: {e}") + + if args.mcp_server in ["valkey", "all"]: + try: + valkey_client = get_valkey_mcp_client(readonly=args.ro) + clients.append(valkey_client) + except Exception as e: + print(f"Warning: Failed to initialize Valkey client: {e}") + + if args.mcp_server in ["memcached", "all"]: + try: + memcached_client = get_memcached_mcp_client(readonly=args.ro) + clients.append(memcached_client) + except Exception as e: + print(f"Warning: Failed to initialize Memcached client: {e}") + + if args.mcp_server in ["aws-docs", "all"]: + try: + aws_docs_client = get_aws_documentation_mcp_client() + clients.append(aws_docs_client) + except Exception as e: + print(f"Warning: Failed to initialize AWS Documentation client: {e}") + + if not clients: + print("Error: No clients were successfully initialized.") + return + + # Execute the prompt with all selected clients + tools = [file_read, file_write] + + # Start all clients and collect their tools + for client in clients: + client.start() + tools.extend(client.list_tools_sync()) + + try: + agent = Agent(tools=tools, model=args.model) + response = agent(args.prompt) + print(response) + except Exception as e: + print(f"Error executing prompt: {e}") + finally: + # Stop all clients + for client in clients: + client.stop(None, None, None) # Provide required arguments for context manager exit method + + +if __name__ == "__main__": + main() diff --git a/03-integrations/elasticache/pyproject.toml b/03-integrations/elasticache/pyproject.toml new file mode 100644 index 00000000..cc2b09b2 --- /dev/null +++ b/03-integrations/elasticache/pyproject.toml @@ -0,0 +1,53 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "elasticache-integration" +version = "0.1.0" +description = "ElastiCache integration for Strands Agent framework" +readme = "README.md" +requires-python = ">=3.10" +license = {file = "LICENSE"} +authors = [ + {name = "seaofawareness", email="utkarshshah@gmail.com"}, +] +dependencies = [ + "strands-agents>=0.1.6", + "strands-agents-tools>=0.1.4", + "python-dotenv>=1.0.0", + "mcp>=0.1.0", + "awslabs.elasticache-mcp-server", + "awslabs.valkey-mcp-server", + "awslabs.memcached-mcp-server", + "awslabs.aws-documentation-mcp-server" +] + +[project.optional-dependencies] +dev = [ + "pytest>=7.0.0", + "pytest-mock>=3.10.0", + "black>=23.0.0", + "isort>=5.12.0", + "flake8>=6.0.0", + "mypy>=1.0.0", +] + +[tool.black] +line-length = 88 +target-version = ["py310"] + +[tool.isort] +profile = "black" +line_length = 88 + +[tool.mypy] +python_version = "3.10" +warn_return_any = true +warn_unused_configs = true +disallow_untyped_defs = true +disallow_incomplete_defs = true + +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = "test_*.py" diff --git a/03-integrations/elasticache/requirements.txt b/03-integrations/elasticache/requirements.txt new file mode 100644 index 00000000..305b7195 --- /dev/null +++ b/03-integrations/elasticache/requirements.txt @@ -0,0 +1,11 @@ +strands-agents>=0.1.6 +strands-agents-tools>0.1.4 +python-dotenv>=1.0.0 +mcp>=0.1.0 +pytest>=7.0.0 +pytest-mock>=3.10.0 + +awslabs.elasticache-mcp-server +awslabs.valkey-mcp-server +awslabs.memcached-mcp-server +awslabs.aws-documentation-mcp-server diff --git a/03-integrations/elasticache/tests/test_main.py b/03-integrations/elasticache/tests/test_main.py new file mode 100644 index 00000000..c2baefca --- /dev/null +++ b/03-integrations/elasticache/tests/test_main.py @@ -0,0 +1,297 @@ +#!/usr/bin/env python3 +import os +import unittest +from unittest.mock import MagicMock, patch + +import pytest + + +class TestElastiCacheIntegration(unittest.TestCase): + """Test cases for ElastiCache integration.""" + + @patch("os.getenv") + def test_get_elasticache_mcp_client(self, mock_getenv): + """Test creating ElastiCache MCP client.""" + # Import here to avoid module-level import issues during testing + from main import get_elasticache_mcp_client + + # Mock environment variables + def mock_env(var, default=None): + env_vars = { + "AWS_REGION": "us-east-1", + "AWS_PROFILE": "test-profile", + "ELASTICACHE_MAX_RETRIES": "5", + "ELASTICACHE_RETRY_MODE": "adaptive", + "ELASTICACHE_CONNECT_TIMEOUT": "10", + "ELASTICACHE_READ_TIMEOUT": "15", + "COST_EXPLORER_MAX_RETRIES": "4", + "COST_EXPLORER_RETRY_MODE": "standard", + "COST_EXPLORER_CONNECT_TIMEOUT": "8", + "COST_EXPLORER_READ_TIMEOUT": "12", + "CLOUDWATCH_MAX_RETRIES": "3", + "CLOUDWATCH_RETRY_MODE": "standard", + "CLOUDWATCH_CONNECT_TIMEOUT": "7", + "CLOUDWATCH_READ_TIMEOUT": "14", + "CLOUDWATCH_LOGS_MAX_RETRIES": "6", + "CLOUDWATCH_LOGS_RETRY_MODE": "adaptive", + "CLOUDWATCH_LOGS_CONNECT_TIMEOUT": "9", + "CLOUDWATCH_LOGS_READ_TIMEOUT": "18", + "FIREHOSE_MAX_RETRIES": "4", + "FIREHOSE_RETRY_MODE": "standard", + "FIREHOSE_CONNECT_TIMEOUT": "6", + "FIREHOSE_READ_TIMEOUT": "11" + } + return env_vars.get(var, default) + + mock_getenv.side_effect = mock_env + + # Mock MCPClient and StdioServerParameters + with patch("main.MCPClient") as mock_mcp_client, \ + patch("main.stdio_client") as mock_stdio_client, \ + patch("main.StdioServerParameters") as mock_params: + + # Set up mocks + mock_params.return_value = "mock_params" + mock_stdio_client.return_value = "mock_stdio_client" + mock_mcp_client.return_value = "mock_client" + + # Call the function + client = get_elasticache_mcp_client() + + # Verify the function called the right methods with the right arguments + mock_params.assert_called_once_with( + command="uvx", + args=["awslabs.elasticache-mcp-server@latest"], + env={ + "AWS_REGION": "us-east-1", + "AWS_PROFILE": "test-profile", + "ELASTICACHE_MAX_RETRIES": "5", + "ELASTICACHE_RETRY_MODE": "adaptive", + "ELASTICACHE_CONNECT_TIMEOUT": "10", + "ELASTICACHE_READ_TIMEOUT": "15", + "COST_EXPLORER_MAX_RETRIES": "4", + "COST_EXPLORER_RETRY_MODE": "standard", + "COST_EXPLORER_CONNECT_TIMEOUT": "8", + "COST_EXPLORER_READ_TIMEOUT": "12", + "CLOUDWATCH_MAX_RETRIES": "3", + "CLOUDWATCH_RETRY_MODE": "standard", + "CLOUDWATCH_CONNECT_TIMEOUT": "7", + "CLOUDWATCH_READ_TIMEOUT": "14", + "CLOUDWATCH_LOGS_MAX_RETRIES": "6", + "CLOUDWATCH_LOGS_RETRY_MODE": "adaptive", + "CLOUDWATCH_LOGS_CONNECT_TIMEOUT": "9", + "CLOUDWATCH_LOGS_READ_TIMEOUT": "18", + "FIREHOSE_MAX_RETRIES": "4", + "FIREHOSE_RETRY_MODE": "standard", + "FIREHOSE_CONNECT_TIMEOUT": "6", + "FIREHOSE_READ_TIMEOUT": "11" + } + ) + mock_stdio_client.assert_called_once_with("mock_params") + mock_mcp_client.assert_called_once() + + # Verify the function returned the expected client + self.assertEqual(client, "mock_client") + + @patch("os.getenv") + def test_get_valkey_mcp_client(self, mock_getenv): + """Test creating Valkey MCP client.""" + # Import here to avoid module-level import issues during testing + from main import get_valkey_mcp_client + + # Mock environment variables + def mock_env(var, default=None): + if var == "VALKEY_HOST": + return "redis.example.com" + elif var == "VALKEY_PORT": + return "6380" + elif var == "VALKEY_USERNAME": + return "user" + elif var == "VALKEY_PWD": + return "password" + elif var == "VALKEY_USE_SSL": + return "True" + elif var == "VALKEY_CLUSTER_MODE": + return "True" + return default + + mock_getenv.side_effect = mock_env + + # Mock MCPClient and StdioServerParameters + with patch("main.MCPClient") as mock_mcp_client, \ + patch("main.stdio_client") as mock_stdio_client, \ + patch("main.StdioServerParameters") as mock_params: + + # Set up mocks + mock_params.return_value = "mock_params" + mock_stdio_client.return_value = "mock_stdio_client" + mock_mcp_client.return_value = "mock_client" + + # Call the function + client = get_valkey_mcp_client() + + # Verify the function called the right methods with the right arguments + mock_params.assert_called_once_with( + command="uvx", + args=["awslabs.valkey-mcp-server@latest"], + env={ + "VALKEY_HOST": "redis.example.com", + "VALKEY_PORT": "6380", + "VALKEY_USERNAME": "user", + "VALKEY_PWD": "password", + "VALKEY_USE_SSL": "true", + "VALKEY_CLUSTER_MODE": "true" + } + ) + mock_stdio_client.assert_called_once_with("mock_params") + mock_mcp_client.assert_called_once() + + # Verify the function returned the expected client + self.assertEqual(client, "mock_client") + + @patch("os.getenv") + def test_get_memcached_mcp_client(self, mock_getenv): + """Test creating Memcached MCP client.""" + # Import here to avoid module-level import issues during testing + from main import get_memcached_mcp_client + + # Mock environment variables + def mock_env(var, default=None): + if var == "MEMCACHED_HOST": + return "memcached.example.com" + elif var == "MEMCACHED_PORT": + return "11212" + return default + + mock_getenv.side_effect = mock_env + + # Mock MCPClient and StdioServerParameters + with patch("main.MCPClient") as mock_mcp_client, \ + patch("main.stdio_client") as mock_stdio_client, \ + patch("main.StdioServerParameters") as mock_params: + + # Set up mocks + mock_params.return_value = "mock_params" + mock_stdio_client.return_value = "mock_stdio_client" + mock_mcp_client.return_value = "mock_client" + + # Call the function + client = get_memcached_mcp_client() + + # Verify the function called the right methods with the right arguments + mock_params.assert_called_once_with( + command="uvx", + args=["awslabs.memcached-mcp-server@latest"], + env={ + "MEMCACHED_HOST": "memcached.example.com", + "MEMCACHED_PORT": "11212" + } + ) + mock_stdio_client.assert_called_once_with("mock_params") + mock_mcp_client.assert_called_once() + + # Verify the function returned the expected client + self.assertEqual(client, "mock_client") + + @patch("os.getenv") + def test_get_aws_documentation_mcp_client(self, mock_getenv): + """Test creating AWS Documentation MCP client.""" + # Import here to avoid module-level import issues during testing + from main import get_aws_documentation_mcp_client + + # Mock environment variables + mock_getenv.return_value = "aws-cn" + + # Mock MCPClient and StdioServerParameters + with patch("main.MCPClient") as mock_mcp_client, \ + patch("main.stdio_client") as mock_stdio_client, \ + patch("main.StdioServerParameters") as mock_params: + + # Set up mocks + mock_params.return_value = "mock_params" + mock_stdio_client.return_value = "mock_stdio_client" + mock_mcp_client.return_value = "mock_client" + + # Call the function + client = get_aws_documentation_mcp_client() + + # Verify the function called the right methods with the right arguments + mock_params.assert_called_once_with( + command="uvx", + args=["awslabs.aws-documentation-mcp-server@latest"], + env={"AWS_DOCUMENTATION_PARTITION": "aws-cn"} + ) + mock_stdio_client.assert_called_once_with("mock_params") + mock_mcp_client.assert_called_once() + + # Verify the function returned the expected client + self.assertEqual(client, "mock_client") + + @patch("main.get_elasticache_mcp_client") + @patch("main.get_valkey_mcp_client") + @patch("main.get_memcached_mcp_client") + @patch("main.get_aws_documentation_mcp_client") + @patch("main.Agent") + def test_main_with_all_engines(self, mock_agent, mock_aws_docs, mock_memcached, mock_valkey, mock_elasticache): + """Test main function with all engines.""" + # Import here to avoid module-level import issues during testing + import main + from argparse import Namespace + + # Mock command line arguments + args = Namespace(prompt="Test prompt", mcp_server="all", ro=False) + with patch("main.argparse.ArgumentParser.parse_args", return_value=args): + + # Set up mocks + mock_elasticache_client = MagicMock() + mock_valkey_client = MagicMock() + mock_memcached_client = MagicMock() + mock_aws_docs_client = MagicMock() + + mock_elasticache.return_value = mock_elasticache_client + mock_valkey.return_value = mock_valkey_client + mock_memcached.return_value = mock_memcached_client + mock_aws_docs.return_value = mock_aws_docs_client + + mock_elasticache_client.list_tools_sync.return_value = ["elasticache_tool"] + mock_valkey_client.list_tools_sync.return_value = ["valkey_tool"] + mock_memcached_client.list_tools_sync.return_value = ["memcached_tool"] + mock_aws_docs_client.list_tools_sync.return_value = ["aws_docs_tool"] + + mock_agent_instance = MagicMock() + mock_agent.return_value = mock_agent_instance + mock_agent_instance.return_value = "Agent response" + + # Call the main function + with patch("main.print") as mock_print: + main.main() + + # Verify the clients were started and stopped + mock_elasticache_client.start.assert_called_once() + mock_valkey_client.start.assert_called_once() + mock_memcached_client.start.assert_called_once() + mock_aws_docs_client.start.assert_called_once() + + mock_elasticache_client.stop.assert_called_once() + mock_valkey_client.stop.assert_called_once() + mock_memcached_client.stop.assert_called_once() + mock_aws_docs_client.stop.assert_called_once() + + # Verify the agent was created with the right tools + mock_agent.assert_called_once() + args, kwargs = mock_agent.call_args + tools = kwargs.get("tools", []) + self.assertIn("elasticache_tool", tools) + self.assertIn("valkey_tool", tools) + self.assertIn("memcached_tool", tools) + self.assertIn("aws_docs_tool", tools) + + # Verify the agent was called with the right prompt + mock_agent_instance.assert_called_once_with("Test prompt") + + # Verify the response was printed + mock_print.assert_called_once_with("Agent response") + + +if __name__ == "__main__": + unittest.main()