#!/usr/bin/env python3
"""
Test suite for API Manager MCP Server
"""
import os
import tempfile
import pytest
import asyncio
from pathlib import Path
from unittest.mock import patch, MagicMock
# Import the main module
import sys
sys.path.insert(0, str(Path(__file__).parent.parent))
from api_manager import APIKeyManager, APIKeyStats
class TestAPIKeyManager:
"""Test cases for APIKeyManager class"""
def setup_method(self):
"""Setup test environment before each test"""
# Create temporary directory for test files
self.temp_dir = tempfile.mkdtemp()
self.test_env_file = Path(self.temp_dir) / ".env"
self.manager = APIKeyManager(str(self.test_env_file))
def teardown_method(self):
"""Cleanup after each test"""
# Clean up temporary files
import shutil
shutil.rmtree(self.temp_dir, ignore_errors=True)
def test_initialization(self):
"""Test APIKeyManager initialization"""
assert self.manager.env_file == self.test_env_file
assert self.test_env_file.exists()
def test_add_initial_header(self):
"""Test initial header creation"""
content = self.test_env_file.read_text()
assert "API Manager Environment File" in content
assert "Created on" in content
def test_load_empty_env(self):
"""Test loading empty environment file"""
env_vars = self.manager.load_env()
assert env_vars == {}
def test_save_and_load_env(self):
"""Test saving and loading environment variables"""
test_vars = {
"OPENAI_API_KEY": "sk-1234567890abcdef",
"ANTHROPIC_API_KEY": "sk-ant-1234567890abcdef"
}
self.manager.save_env(test_vars)
loaded_vars = self.manager.load_env()
assert loaded_vars == test_vars
def test_get_category(self):
"""Test category determination"""
test_cases = [
("OPENAI_API_KEY", "OpenAI"),
("ANTHROPIC_API_KEY", "Anthropic (Claude)"),
("DISCORD_BOT_TOKEN", "Discord"),
("UNKNOWN_KEY", "Other"),
("SUPABASE_API_KEY", "Supabase"),
("GITHUB_TOKEN", "GitHub")
]
for key, expected_category in test_cases:
assert self.manager.get_category(key) == expected_category
def test_validate_key_name(self):
"""Test key name validation"""
valid_names = ["OPENAI_API_KEY", "test_key", "API_KEY_123"]
invalid_names = ["", "key/with/slash", "key*with*asterisk", "key:with:colon"]
for name in valid_names:
assert self.manager.validate_key_name(name), f"'{name}' should be valid"
for name in invalid_names:
assert not self.manager.validate_key_name(name), f"'{name}' should be invalid"
def test_validate_key_value(self):
"""Test key value validation"""
valid_values = [
"sk-1234567890abcdef",
"very-long-api-key-value-123",
"AIzaSyDummyGoogleAPIKey123456789"
]
invalid_values = [
"",
"short",
"test",
"example",
"your_key_here"
]
for value in valid_values:
assert self.manager.validate_key_value(value), f"'{value}' should be valid"
for value in invalid_values:
assert not self.manager.validate_key_value(value), f"'{value}' should be invalid"
def test_mask_value(self):
"""Test value masking"""
test_cases = [
("short", "*****"),
("mediumkey", "me****ey"),
("very-long-api-key-value", "very************alue")
]
for value, expected in test_cases:
assert self.manager.mask_value(value) == expected
def test_get_stats(self):
"""Test statistics generation"""
test_vars = {
"OPENAI_API_KEY": "sk-1234567890abcdef",
"ANTHROPIC_API_KEY": "sk-ant-1234567890abcdef",
"DISCORD_BOT_TOKEN": "discord-token-123",
"GITHUB_TOKEN": "github-token-456"
}
self.manager.save_env(test_vars)
stats = self.manager.get_stats()
assert isinstance(stats, APIKeyStats)
assert stats.total_keys == 4
assert "OpenAI" in stats.categories
assert "Anthropic (Claude)" in stats.categories
assert "Discord" in stats.categories
assert "GitHub" in stats.categories
def test_create_backup(self):
"""Test backup creation"""
# Add some test data
test_vars = {"TEST_KEY": "test_value"}
self.manager.save_env(test_vars)
# Create backup
backup_file = self.manager.create_backup()
assert Path(backup_file).exists()
assert "backup_" in backup_file
def test_export_keys_json(self):
"""Test JSON export"""
test_vars = {
"OPENAI_API_KEY": "sk-1234567890abcdef",
"DISCORD_TOKEN": "discord-token-123"
}
self.manager.save_env(test_vars)
# Test masked export
json_export = self.manager.export_keys("json", include_values=False)
assert "sk-1***" in json_export
assert "discord-token-123" not in json_export
# Test full export
json_export_full = self.manager.export_keys("json", include_values=True)
assert "sk-1234567890abcdef" in json_export_full
assert "discord-token-123" in json_export_full
def test_export_keys_env(self):
"""Test ENV format export"""
test_vars = {
"OPENAI_API_KEY": "sk-1234567890abcdef",
"DISCORD_TOKEN": "discord-token-123"
}
self.manager.save_env(test_vars)
# Test masked export
env_export = self.manager.export_keys("env", include_values=False)
assert "OPENAI_API_KEY=" in env_export
assert "sk-1234567890abcdef" not in env_export
def test_export_invalid_format(self):
"""Test export with invalid format"""
with pytest.raises(ValueError):
self.manager.export_keys("invalid_format")
def test_atomic_write(self):
"""Test atomic file writing"""
# Add initial data
initial_vars = {"KEY1": "value1"}
self.manager.save_env(initial_vars)
# Simulate concurrent access
updated_vars = {"KEY1": "updated_value1", "KEY2": "value2"}
# This should work without corruption
self.manager.save_env(updated_vars)
loaded_vars = self.manager.load_env()
assert loaded_vars == updated_vars
def test_preserve_comments(self):
"""Test that comments are preserved when updating"""
# Manually add a file with comments
content = """# Test comment
KEY1=value1
# Another comment
KEY2=value2
"""
self.test_env_file.write_text(content)
# Update with new values
updated_vars = {"KEY1": "new_value1", "KEY2": "value2", "KEY3": "value3"}
self.manager.save_env(updated_vars)
# Check that comments are preserved
final_content = self.test_env_file.read_text()
assert "# Test comment" in final_content
assert "# Another comment" in final_content
assert "KEY1=new_value1" in final_content
def test_file_permissions(self):
"""Test that file permissions are set correctly"""
# Add some data to trigger file creation
self.manager.save_env({"TEST_KEY": "test_value"})
# Check file permissions (should be 600)
file_mode = self.test_env_file.stat().st_mode & 0o777
assert file_mode == 0o600
def test_malformed_env_file(self):
"""Test handling of malformed environment file"""
# Create malformed file
malformed_content = """
KEY1=value1
MALFORMED_LINE_WITHOUT_EQUALS
KEY2=value2
=VALUE_WITHOUT_KEY
KEY3=value3=with=extra=equals
"""
self.test_env_file.write_text(malformed_content)
# Should still load valid entries
env_vars = self.manager.load_env()
assert "KEY1" in env_vars
assert "KEY2" in env_vars
assert "KEY3" in env_vars
assert env_vars["KEY3"] == "value3=with=extra=equals"
class TestAsyncOperations:
"""Test async operations and MCP integration"""
@pytest.fixture
def event_loop(self):
"""Create event loop for async tests"""
loop = asyncio.new_event_loop()
yield loop
loop.close()
@pytest.mark.asyncio
async def test_list_tools(self):
"""Test tool listing"""
# Import required modules for testing
from api_manager import app
tools = await app._tool_handlers["list_tools"]()
assert len(tools) > 0
tool_names = [tool.name for tool in tools]
expected_tools = [
"list_api_keys",
"get_api_key",
"add_api_key",
"delete_api_key",
"search_api_keys",
"list_categories",
"backup_env_file",
"get_stats",
"export_keys",
"help"
]
for expected_tool in expected_tools:
assert expected_tool in tool_names
@pytest.mark.asyncio
async def test_call_tool_help(self):
"""Test help tool call"""
from api_manager import call_tool
result = await call_tool("help", {})
assert len(result) == 1
assert "API Manager MCP Server" in result[0].text
@pytest.mark.asyncio
async def test_call_tool_unknown(self):
"""Test unknown tool call"""
from api_manager import call_tool
result = await call_tool("unknown_tool", {})
assert len(result) == 1
assert "Unknown tool" in result[0].text
if __name__ == "__main__":
pytest.main([__file__])
