290 lines
8.8 KiB
Python
Executable File
290 lines
8.8 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""
|
||
Resource Limits Enforcement Test Script
|
||
|
||
Tests that container resource limits are properly applied and enforced,
|
||
preventing resource exhaustion attacks and ensuring fair resource allocation.
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import asyncio
|
||
import json
|
||
import time
|
||
from pathlib import Path
|
||
|
||
# Add session-manager to path for imports
|
||
sys.path.insert(0, str(Path(__file__).parent))
|
||
|
||
from resource_manager import (
|
||
get_resource_limits,
|
||
check_system_resources,
|
||
should_throttle_sessions,
|
||
ResourceLimits,
|
||
ResourceValidator,
|
||
ResourceMonitor,
|
||
)
|
||
|
||
# Set up logging
|
||
import logging
|
||
|
||
logging.basicConfig(level=logging.INFO)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
async def test_resource_limits_configuration():
|
||
"""Test resource limits configuration and validation."""
|
||
print("🧪 Testing Resource Limits Configuration")
|
||
print("=" * 50)
|
||
|
||
# Test default configuration
|
||
print("1️⃣ Testing default configuration...")
|
||
try:
|
||
limits = get_resource_limits()
|
||
print(
|
||
f"✅ Default limits loaded: memory={limits.memory_limit}, cpu_quota={limits.cpu_quota}"
|
||
)
|
||
|
||
# Validate limits
|
||
valid, message = limits.validate()
|
||
if valid:
|
||
print("✅ Default limits validation passed")
|
||
else:
|
||
print(f"❌ Default limits validation failed: {message}")
|
||
return False
|
||
except Exception as e:
|
||
print(f"❌ Failed to load default configuration: {e}")
|
||
return False
|
||
|
||
# Test custom configuration
|
||
print("\n2️⃣ Testing custom configuration...")
|
||
test_configs = [
|
||
{"memory_limit": "2g", "cpu_quota": 50000, "cpu_period": 100000},
|
||
{"memory_limit": "512m", "cpu_quota": 25000, "cpu_period": 100000},
|
||
{"memory_limit": "8g", "cpu_quota": 200000, "cpu_period": 100000},
|
||
]
|
||
|
||
for config in test_configs:
|
||
try:
|
||
valid, message, limits = ResourceValidator.validate_resource_config(config)
|
||
if valid:
|
||
print(f"✅ Config {config} validated successfully")
|
||
else:
|
||
print(f"❌ Config {config} validation failed: {message}")
|
||
return False
|
||
except Exception as e:
|
||
print(f"❌ Config {config} validation error: {e}")
|
||
return False
|
||
|
||
# Test invalid configurations
|
||
print("\n3️⃣ Testing invalid configurations...")
|
||
invalid_configs = [
|
||
{
|
||
"memory_limit": "0g",
|
||
"cpu_quota": 100000,
|
||
"cpu_period": 100000,
|
||
}, # Zero memory
|
||
{
|
||
"memory_limit": "1g",
|
||
"cpu_quota": 200000,
|
||
"cpu_period": 100000,
|
||
}, # CPU quota > period
|
||
{
|
||
"memory_limit": "1g",
|
||
"cpu_quota": -1000,
|
||
"cpu_period": 100000,
|
||
}, # Negative CPU
|
||
]
|
||
|
||
for config in invalid_configs:
|
||
valid, message, limits = ResourceValidator.validate_resource_config(config)
|
||
if not valid:
|
||
print(f"✅ Correctly rejected invalid config {config}: {message}")
|
||
else:
|
||
print(f"❌ Incorrectly accepted invalid config {config}")
|
||
return False
|
||
|
||
return True
|
||
|
||
|
||
async def test_resource_monitoring():
|
||
"""Test system resource monitoring."""
|
||
print("\n🖥️ Testing System Resource Monitoring")
|
||
print("=" * 50)
|
||
|
||
# Test resource monitoring
|
||
print("1️⃣ Testing system resource monitoring...")
|
||
try:
|
||
resource_check = check_system_resources()
|
||
print("✅ System resource check completed")
|
||
|
||
if isinstance(resource_check, dict):
|
||
system_resources = resource_check.get("system_resources", {})
|
||
alerts = resource_check.get("alerts", [])
|
||
|
||
print(f"📊 System resources: {len(system_resources)} metrics collected")
|
||
print(f"🚨 Resource alerts: {len(alerts)} detected")
|
||
|
||
for alert in alerts:
|
||
print(
|
||
f" {alert.get('level', 'unknown').upper()}: {alert.get('message', 'No message')}"
|
||
)
|
||
|
||
# Test throttling logic
|
||
should_throttle, reason = should_throttle_sessions()
|
||
print(f"🎛️ Session throttling: {'YES' if should_throttle else 'NO'} - {reason}")
|
||
|
||
except Exception as e:
|
||
print(f"❌ Resource monitoring test failed: {e}")
|
||
return False
|
||
|
||
return True
|
||
|
||
|
||
async def test_memory_limit_parsing():
|
||
"""Test memory limit parsing functionality."""
|
||
print("\n💾 Testing Memory Limit Parsing")
|
||
print("=" * 50)
|
||
|
||
test_cases = [
|
||
("4g", (4 * 1024**3, "GB")),
|
||
("512m", (512 * 1024**2, "MB")),
|
||
("256k", (256 * 1024, "KB")),
|
||
("1073741824", (1073741824, "bytes")), # 1GB in bytes
|
||
]
|
||
|
||
for memory_str, expected in test_cases:
|
||
try:
|
||
bytes_val, unit = ResourceValidator.parse_memory_limit(memory_str)
|
||
if bytes_val == expected[0] and unit == expected[1]:
|
||
print(f"✅ Parsed {memory_str} -> {bytes_val} {unit}")
|
||
else:
|
||
print(
|
||
f"❌ Failed to parse {memory_str}: got {bytes_val} {unit}, expected {expected}"
|
||
)
|
||
return False
|
||
except Exception as e:
|
||
print(f"❌ Error parsing {memory_str}: {e}")
|
||
return False
|
||
|
||
return True
|
||
|
||
|
||
async def test_docker_limits_conversion():
|
||
"""Test Docker limits conversion."""
|
||
print("\n🐳 Testing Docker Limits Conversion")
|
||
print("=" * 50)
|
||
|
||
limits = ResourceLimits(
|
||
memory_limit="2g",
|
||
cpu_quota=75000,
|
||
cpu_period=100000,
|
||
)
|
||
|
||
docker_limits = limits.to_docker_limits()
|
||
expected = {
|
||
"mem_limit": "2g",
|
||
"cpu_quota": 75000,
|
||
"cpu_period": 100000,
|
||
}
|
||
|
||
if docker_limits == expected:
|
||
print("✅ Docker limits conversion correct")
|
||
return True
|
||
else:
|
||
print(f"❌ Docker limits conversion failed: {docker_limits} != {expected}")
|
||
return False
|
||
|
||
|
||
async def test_environment_variables():
|
||
"""Test environment variable configuration."""
|
||
print("\n🌍 Testing Environment Variable Configuration")
|
||
print("=" * 50)
|
||
|
||
# Save original environment
|
||
original_env = {}
|
||
test_vars = [
|
||
"CONTAINER_MEMORY_LIMIT",
|
||
"CONTAINER_CPU_QUOTA",
|
||
"CONTAINER_CPU_PERIOD",
|
||
"MAX_CONCURRENT_SESSIONS",
|
||
"MEMORY_WARNING_THRESHOLD",
|
||
"CPU_WARNING_THRESHOLD",
|
||
]
|
||
|
||
for var in test_vars:
|
||
original_env[var] = os.environ.get(var)
|
||
|
||
try:
|
||
# Test custom environment configuration
|
||
os.environ["CONTAINER_MEMORY_LIMIT"] = "3g"
|
||
os.environ["CONTAINER_CPU_QUOTA"] = "80000"
|
||
os.environ["CONTAINER_CPU_PERIOD"] = "100000"
|
||
os.environ["MAX_CONCURRENT_SESSIONS"] = "5"
|
||
|
||
# Force reload of configuration
|
||
import importlib
|
||
import resource_manager
|
||
|
||
importlib.reload(resource_manager)
|
||
|
||
limits = resource_manager.get_resource_limits()
|
||
|
||
if limits.memory_limit == "3g" and limits.cpu_quota == 80000:
|
||
print("✅ Environment variable configuration working")
|
||
return True
|
||
else:
|
||
print(f"❌ Environment variable configuration failed: got {limits}")
|
||
return False
|
||
|
||
finally:
|
||
# Restore original environment
|
||
for var, value in original_env.items():
|
||
if value is not None:
|
||
os.environ[var] = value
|
||
elif var in os.environ:
|
||
del os.environ[var]
|
||
|
||
|
||
async def run_all_tests():
|
||
"""Run all resource limit tests."""
|
||
print("🚀 Resource Limits Enforcement Test Suite")
|
||
print("=" * 60)
|
||
|
||
tests = [
|
||
("Resource Limits Configuration", test_resource_limits_configuration),
|
||
("System Resource Monitoring", test_resource_monitoring),
|
||
("Memory Limit Parsing", test_memory_limit_parsing),
|
||
("Docker Limits Conversion", test_docker_limits_conversion),
|
||
("Environment Variables", test_environment_variables),
|
||
]
|
||
|
||
results = []
|
||
for test_name, test_func in tests:
|
||
print(f"\n{'=' * 20} {test_name} {'=' * 20}")
|
||
try:
|
||
result = await test_func()
|
||
results.append(result)
|
||
status = "✅ PASSED" if result else "❌ FAILED"
|
||
print(f"\n{status}: {test_name}")
|
||
except Exception as e:
|
||
print(f"\n❌ ERROR in {test_name}: {e}")
|
||
results.append(False)
|
||
|
||
# Summary
|
||
print(f"\n{'=' * 60}")
|
||
passed = sum(results)
|
||
total = len(results)
|
||
print(f"📊 Test Results: {passed}/{total} tests passed")
|
||
|
||
if passed == total:
|
||
print("🎉 All resource limit tests completed successfully!")
|
||
print("💪 Container resource limits are properly configured and enforced.")
|
||
else:
|
||
print("⚠️ Some tests failed. Check the output above for details.")
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(run_all_tests())
|