rustytorch/crates/training/rtx-federated/tests/integration_test.rs

//! Integration tests for rtx-federated
//!
//! These tests follow strict TDD methodology with RED-GREEN-REFACTOR cycle
#![cfg(feature = "disabled_tests")]

use rtx_federated::{
    aggregation::{AggregationAlgorithm, FedAvg, ModelUpdate},
    error::Result,
    infrastructure::client_manager::{ClientManager, protocol::ClientConnection},
};
use uuid::Uuid;

#[tokio::test]
async fn test_client_manager_creation() -> Result<()> {
    // Test that client manager can be created
    let _manager = ClientManager::new().await?;
    Ok(())
}

#[tokio::test]
async fn test_client_registration() -> Result<()> {
    // RED: Test client registration flow
    let client_manager = ClientManager::new().await?;

    // Generate a client ID
    let client_id = Uuid::new_v4();

    // Create a mock connection (we'll need to implement this properly)
    let connection = create_test_connection();

    // Register the client
    client_manager
        .register_client(client_id, connection)
        .await?;

    // Verify client is registered
    assert_eq!(client_manager.connected_clients(), 1);

    Ok(())
}

#[tokio::test]
async fn test_federated_averaging() -> Result<()> {
    // Test FedAvg aggregation
    let fedavg = FedAvg::new(Some(0.9), true).await?;

    // Create some test model updates
    let updates = create_test_model_updates(3);

    // Perform aggregation
    let aggregated = fedavg.aggregate(&updates).await?;

    // Verify aggregation produced a result
    assert!(!aggregated.parameters.is_empty());

    Ok(())
}

#[tokio::test]
async fn test_privacy_budget_management() -> Result<()> {
    // Test privacy budget tracking
    use rtx_federated::privacy::{PrivacyAccountant, PrivacyBudget};

    let total_budget = PrivacyBudget {
        epsilon: 10.0,
        delta: 1e-5,
    };
    let mut accountant = PrivacyAccountant::new(total_budget);

    // Consume some privacy budget
    accountant.consume(1.0, 1e-6, "test operation".to_string())?;

    // Verify budget was consumed
    let remaining = accountant.remaining_budget();
    assert_eq!(remaining.epsilon, 9.0);

    Ok(())
}

#[tokio::test]
async fn test_byzantine_detection() -> Result<()> {
    // Test byzantine fault detection using Krum
    use rtx_federated::byzantine::{ByzantineRobust, Krum};

    let krum = Krum::new(0.2).await?;

    // Create updates with one malicious client
    let mut updates = create_test_model_updates(5);
    corrupt_update(&mut updates[2]); // Make one update malicious

    // Filter updates (Krum should detect and filter the malicious one)
    let filtered_updates = krum.filter_updates(&updates).await?;

    // Verify filtering worked (should have fewer updates)
    assert!(filtered_updates.len() < updates.len());

    Ok(())
}

// Helper functions for testing

fn create_test_connection() -> ClientConnection {
    use rtx_federated::infrastructure::client_manager::protocol::ConnectionQuality;
    use std::net::SocketAddr;
    use std::sync::Arc;
    use std::sync::atomic::AtomicU64;

    ClientConnection {
        client_id: Uuid::new_v4(),
        endpoint: "127.0.0.1:8080".parse::<SocketAddr>().unwrap(),
        connected_at: chrono::Utc::now(),
        last_heartbeat: chrono::Utc::now(),
        bytes_sent: Arc::new(AtomicU64::new(0)),
        bytes_received: Arc::new(AtomicU64::new(0)),
        round_trip_time_ms: Arc::new(AtomicU64::new(0)),
        connection_quality: ConnectionQuality::default(),
        auth_token: None,
        stream: None,
    }
}

fn create_test_model_updates(count: usize) -> Vec<ModelUpdate> {
    use std::collections::HashMap;

    (0..count)
        .map(|i| {
            let mut update = ModelUpdate::new(Uuid::new_v4());
            update.add_parameter(
                "layer1.weight",
                vec![0.1 * (i as f64), 0.2 * (i as f64), 0.3 * (i as f64)],
            );
            update.add_parameter("layer1.bias", vec![0.01 * (i as f64)]);
            update.sample_count = 100 + (i * 10);
            update.loss = 0.5;
            update.accuracy = 0.8;
            update
        })
        .collect()
}

fn corrupt_update(update: &mut ModelUpdate) {
    // Corrupt the parameters to simulate a byzantine attack
    for (_, values) in update.parameters.iter_mut() {
        for v in values.iter_mut() {
            *v *= 100.0; // Amplify values significantly
        }
    }
}

#[tokio::test]
async fn test_basic_aggregation() -> Result<()> {
    // Test basic aggregation without encryption
    let fedavg = FedAvg::new(None, false).await?;

    // Create test updates
    let updates = create_test_model_updates(3);

    // Perform aggregation
    let result = fedavg.aggregate(&updates).await?;

    // Verify result is valid
    assert!(!result.parameters.is_empty());
    assert!(result.sample_count > 0);

    Ok(())
}

#[tokio::test]
async fn test_differential_privacy() -> Result<()> {
    // Test differential privacy mechanism
    use rtx_federated::privacy::{DifferentialPrivacy, PrivacyMechanism};

    let dp = DifferentialPrivacy::new(1.0, 1e-5).await?;

    // Create a model update
    let update = create_test_model_updates(1).pop().unwrap();

    // Apply privacy (which includes adding noise)
    let private_update = dp.apply_privacy(&update).await?;

    // Verify parameters still exist
    assert!(private_update.parameters.values().all(|v| !v.is_empty()));

    Ok(())
}

#[tokio::test]
async fn test_client_registration_and_stats() -> Result<()> {
    // Test client registration and statistics
    let manager = ClientManager::new().await?;

    // Register multiple clients
    for i in 0..5 {
        let connection = create_test_connection();
        manager.register_client(Uuid::new_v4(), connection).await?;
    }

    // Verify registration worked
    let total_clients = manager.connected_clients();
    assert_eq!(total_clients, 5);

    Ok(())
}

// Helper function to create test connections
fn create_multiple_test_connections(count: usize) -> Vec<ClientConnection> {
    use rtx_federated::infrastructure::client_manager::protocol::ConnectionQuality;
    use std::net::SocketAddr;
    use std::sync::Arc;
    use std::sync::atomic::AtomicU64;

    (0..count)
        .map(|_i| ClientConnection {
            client_id: Uuid::new_v4(),
            endpoint: "127.0.0.1:8080".parse::<SocketAddr>().unwrap(),
            connected_at: chrono::Utc::now(),
            last_heartbeat: chrono::Utc::now(),
            bytes_sent: Arc::new(AtomicU64::new(0)),
            bytes_received: Arc::new(AtomicU64::new(0)),
            round_trip_time_ms: Arc::new(AtomicU64::new(0)),
            connection_quality: ConnectionQuality::default(),
            auth_token: None,
            stream: None,
        })
        .collect()
}