Spaces:
Sleeping
Sleeping
#!/usr/bin/env python3 | |
""" | |
Simple test script for the federated learning implementation | |
""" | |
import sys | |
import time | |
import subprocess | |
import threading | |
import os | |
from pathlib import Path | |
import logging | |
import yaml | |
# Set up debug logging for the test | |
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s %(levelname)s %(message)s') | |
# Add src to path | |
sys.path.append(str(Path(__file__).parent / "src")) | |
def load_client_config(): | |
config_path = Path(__file__).parent / "config" / "client_config.yaml" | |
with open(config_path, 'r') as f: | |
full_config = yaml.safe_load(f) | |
return full_config | |
def test_basic_functionality(): | |
"""Test basic federated learning functionality""" | |
print("Testing FinFedRAG Basic Functionality") | |
print("=" * 50) | |
# Test 1: Import all modules | |
print("Test 1: Testing imports...") | |
try: | |
from src.server.coordinator import FederatedCoordinator | |
from src.client.model import FederatedClient | |
from src.api.server import FederatedAPI | |
from src.api.client import FederatedHTTPClient | |
print("✓ All imports successful") | |
logging.debug("All modules imported successfully.") | |
except ImportError as e: | |
print(f"✗ Import failed: {e}") | |
logging.error(f"Import failed: {e}") | |
return False | |
# Test 2: Create coordinator | |
print("\nTest 2: Testing coordinator creation...") | |
try: | |
config = { | |
'server': { | |
'federated': {'min_clients': 2, 'rounds': 3}, | |
'api': {'host': 'localhost', 'port': 8081}, | |
'aggregation': {'method': 'fedavg', 'weighted': True} | |
}, | |
'model': {'input_dim': 32}, | |
'training': {'learning_rate': 0.001} | |
} | |
logging.debug(f"Coordinator test config: {config}") | |
coordinator = FederatedCoordinator(config) | |
print("✓ Coordinator created successfully") | |
logging.debug("Coordinator created successfully.") | |
except Exception as e: | |
print(f"✗ Coordinator creation failed: {e}") | |
logging.error(f"Coordinator creation failed: {e}") | |
return False | |
# Test 3: Create client | |
print("\nTest 3: Testing client creation...") | |
try: | |
client_config = load_client_config() | |
logging.debug(f"Client test config: {client_config}") | |
client = FederatedClient("test_client", client_config) | |
print("✓ Client created successfully") | |
logging.debug("Client created successfully.") | |
except Exception as e: | |
print(f"✗ Client creation failed: {e}") | |
logging.error(f"Client creation failed: {e}") | |
return False | |
# Test 4: Test HTTP client | |
print("\nTest 4: Testing HTTP client...") | |
try: | |
http_client = FederatedHTTPClient('http://localhost:8081', 'test_client') | |
print("✓ HTTP client created successfully") | |
logging.debug("HTTP client created successfully.") | |
except Exception as e: | |
print(f"✗ HTTP client creation failed: {e}") | |
logging.error(f"HTTP client creation failed: {e}") | |
return False | |
print("\n" + "=" * 50) | |
print("All basic functionality tests passed!") | |
logging.debug("All basic functionality tests passed.") | |
return True | |
def run_integration_test(): | |
"""Run a quick integration test""" | |
print("\nRunning Integration Test") | |
print("=" * 50) | |
# This would start a server and client in separate processes | |
# For now, just validate the configuration files | |
config_dir = Path("config") | |
# Test server config | |
server_config = config_dir / "server_config.yaml" | |
if server_config.exists(): | |
print("✓ Server config exists") | |
logging.debug("Server config exists.") | |
else: | |
print("✗ Server config missing") | |
logging.error("Server config missing.") | |
return False | |
# Test client config | |
client_config = config_dir / "client_config.yaml" | |
if client_config.exists(): | |
print("✓ Client config exists") | |
logging.debug("Client config exists.") | |
else: | |
print("✗ Client config missing") | |
logging.error("Client config missing.") | |
return False | |
print("✓ Configuration files are present") | |
print("✓ Integration test setup complete") | |
logging.debug("Integration test setup complete.") | |
return True | |
if __name__ == "__main__": | |
print("FinFedRAG Test Suite") | |
print("=" * 50) | |
# Change to project directory | |
os.chdir(Path(__file__).parent) | |
success = True | |
# Run basic functionality tests | |
if not test_basic_functionality(): | |
success = False | |
# Run integration tests | |
if not run_integration_test(): | |
success = False | |
print("\n" + "=" * 50) | |
if success: | |
print("🎉 All tests passed!") | |
print("\nTo run the system:") | |
print("1. Start server: python -m src.main --mode server --config config/server_config.yaml") | |
print("2. Start client: python -m src.main --mode client --config config/client_config.yaml") | |
else: | |
print("❌ Some tests failed!") | |
sys.exit(1) | |