heartbeat/test_core_plugins.py

#!/usr/bin/env python3
"""
Test script for all monitoring plugins.

Tests all available plugins including the new ones:
- memory_monitor
- disk_monitor
- network_monitor
- filesystem_info
"""

import asyncio
import sys
import os
import logging

# Add parent directory to path so we can import hbd
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from hbd.plugin import PluginLoader

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)

logger = logging.getLogger(__name__)


def format_bytes(bytes_val):
    """Format bytes into human readable format."""
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if bytes_val < 1024.0:
            return f"{bytes_val:.2f} {unit}"
        bytes_val /= 1024.0
    return f"{bytes_val:.2f} PB"


def print_plugin_data(plugin_name, data, indent=2):
    """Pretty print plugin data."""
    prefix = " " * indent
    
    if isinstance(data, dict):
        for key, value in data.items():
            if isinstance(value, dict):
                print(f"{prefix}{key}:")
                print_plugin_data(plugin_name, value, indent + 2)
            elif isinstance(value, list):
                print(f"{prefix}{key}: [{len(value)} items]")
                if len(value) <= 5:  # Only show small lists
                    for item in value:
                        if isinstance(item, dict):
                            print_plugin_data(plugin_name, item, indent + 2)
                        else:
                            print(f"{prefix}  - {item}")
            else:
                # Format output based on key name for better readability
                if '_bytes' in key or key.endswith('_sent') or key.endswith('_recv') or 'memory_' in key or 'swap_' in key:
                    if isinstance(value, (int, float)) and value > 1024:
                        print(f"{prefix}{key}: {format_bytes(value)} ({value:,})")
                    else:
                        print(f"{prefix}{key}: {value}")
                elif 'percent' in key:
                    print(f"{prefix}{key}: {value:.1f}%")
                elif isinstance(value, float):
                    print(f"{prefix}{key}: {value:.2f}")
                elif isinstance(value, int) and value > 1000:
                    print(f"{prefix}{key}: {value:,}")
                else:
                    print(f"{prefix}{key}: {value}")
    else:
        print(f"{prefix}{data}")


async def main():
    """Main test function."""
    print("="*60)
    print("Plugin System Test Suite")
    print("="*60)
    
    # Load all available plugins using the plugin loader
    from hbd.plugin import PluginRegistry, PluginLoader
    from pathlib import Path
    
    registry = PluginRegistry()
    loader = PluginLoader(registry)
    
    plugin_dir = Path(__file__).parent / "hbd" / "plugins"
    if not plugin_dir.exists():
        print(f"✗ Plugin directory not found: {plugin_dir}")
        return 1
    
    # Load plugins from directory
    count = await loader.load_from_directory(plugin_dir, {})
    
    print(f"\nLoaded {count} plugins:")
    plugins = registry.get_all()
    for plugin in plugins:
        print(f"  - {plugin.name}: {plugin.__class__.__doc__.split('.')[0] if plugin.__class__.__doc__ else 'No description'}")
    
    # Test each plugin
    results = {}
    for plugin in plugins:
        # Skip nagios_runner as it needs specific configuration
        if plugin.name == 'nagios_runner':
            print(f"\n{'='*60}")
            print(f"Skipping: {plugin.name} (requires specific configuration)")
            print(f"{'='*60}")
            results[plugin.name] = True  # Mark as success since it loaded OK
            continue
        
        print(f"\n{'='*60}")
        print(f"Testing: {plugin.name}")
        print(f"{'='*60}")
        
        try:
            # Collect data
            data = await plugin.collect()
            if data:
                if 'error' in data:
                    print(f"✗ Collection error: {data['error']}")
                    results[plugin.name] = False
                else:
                    print(f"✓ Data collected: {len(data)} top-level fields")
                    print_plugin_data(plugin.name, data)
                    results[plugin.name] = True
            else:
                print(f"⚠ No data collected")
                results[plugin.name] = False
        except Exception as e:
            print(f"✗ Failed to collect data: {e}")
            import traceback
            traceback.print_exc()
            results[plugin.name] = False
    
    # Summary
    print(f"\n{'='*60}")
    print("Test Summary")
    print(f"{'='*60}")
    
    success_count = sum(1 for v in results.values() if v)
    total_count = len(results)
    
    print(f"\nResults: {success_count}/{total_count} plugins successful")
    for name, success in results.items():
        status = "✓" if success else "✗"
        print(f"  {status} {name}")
    
    if success_count == total_count:
        print("\n🎉 All plugins passed!")
        return 0
    else:
        print(f"\n⚠ {total_count - success_count} plugin(s) failed")
        return 1


if __name__ == '__main__':
    exit_code = asyncio.run(main())
    sys.exit(exit_code)