rustytorch/integration_tests/src/common.rs

/*!
Common utilities and infrastructure for integration tests.

This module provides shared functionality used across all integration test categories,
including test environment setup, resource management, and utility functions.
*/
#![cfg(feature = "integration-tests")]
/*!
*/

use anyhow::Result;
use std::time::{Duration, Instant};
use tracing::{info, warn, error, debug};
use tokio::time::timeout;
use uuid::Uuid;
use sysinfo::SystemExt;

// Re-export stubs for easy access from test files
pub use crate::stubs::*;

/// Test execution context with resource tracking and cleanup
pub struct TestContext {
    pub test_id: Uuid,
    pub start_time: Instant,
    pub allocated_resources: Vec<AllocatedResource>,
    pub temp_files: Vec<std::path::PathBuf>,
    pub spawned_processes: Vec<tokio::process::Child>,
}

impl Default for TestContext {
    fn default() -> Self {
        Self::new()
    }
}

impl TestContext {
    pub fn new() -> Self {
        Self {
            test_id: Uuid::new_v4(),
            start_time: Instant::now(),
            allocated_resources: Vec::new(),
            temp_files: Vec::new(),
            spawned_processes: Vec::new(),
        }
    }

    pub fn elapsed(&self) -> Duration {
        self.start_time.elapsed()
    }

    pub fn add_resource(&mut self, resource: AllocatedResource) {
        self.allocated_resources.push(resource);
    }

    pub fn add_temp_file(&mut self, path: std::path::PathBuf) {
        self.temp_files.push(path);
    }

    pub fn add_process(&mut self, child: tokio::process::Child) {
        self.spawned_processes.push(child);
    }

    /// Clean up all allocated resources
    pub async fn cleanup(&mut self) -> Result<()> {
        debug!("Cleaning up test context {}", self.test_id);

        // Kill spawned processes
        for mut child in self.spawned_processes.drain(..) {
            if let Err(e) = child.kill().await {
                warn!("Failed to kill process: {}", e);
            }
        }

        // Clean up temporary files
        for path in self.temp_files.drain(..) {
            if let Err(e) = std::fs::remove_file(&path) {
                warn!("Failed to remove temp file {:?}: {}", path, e);
            }
        }

        // Clean up allocated resources
        for resource in self.allocated_resources.drain(..) {
            if let Err(e) = resource.cleanup().await {
                warn!("Failed to cleanup resource: {}", e);
            }
        }

        debug!("Test context cleanup completed in {:?}", self.elapsed());
        Ok(())
    }
}

impl Drop for TestContext {
    fn drop(&mut self) {
        if !self.allocated_resources.is_empty() || !self.temp_files.is_empty() || !self.spawned_processes.is_empty() {
            warn!("TestContext dropped without cleanup! Resources may leak.");
        }
    }
}

/// Represents an allocated resource that needs cleanup
#[derive(Debug)]
pub enum AllocatedResource {
    MemoryBuffer { size_bytes: usize, device_id: usize },
    GpuContext { device_id: usize },
    NetworkPort { port: u16 },
    DatabaseConnection { connection_id: String },
    Container { container_id: String },
    TempDirectory { path: std::path::PathBuf },
}

impl AllocatedResource {
    pub async fn cleanup(self) -> Result<()> {
        match self {
            Self::MemoryBuffer { size_bytes, device_id } => {
                debug!("Cleaning up memory buffer: {} bytes on device {}", size_bytes, device_id);
                // Actual memory cleanup would happen through the allocator
                Ok(())
            }
            Self::GpuContext { device_id } => {
                debug!("Cleaning up GPU context on device {}", device_id);
                // GPU context cleanup
                Ok(())
            }
            Self::NetworkPort { port } => {
                debug!("Released network port {}", port);
                Ok(())
            }
            Self::DatabaseConnection { connection_id } => {
                debug!("Closing database connection {}", connection_id);
                Ok(())
            }
            Self::Container { container_id } => {
                debug!("Stopping container {}", container_id);
                let output = tokio::process::Command::new("docker")
                    .args(["stop", &container_id])
                    .output()
                    .await?;

                if !output.status.success() {
                    warn!("Failed to stop container {}: {}",
                          container_id, String::from_utf8_lossy(&output.stderr));
                }

                // Remove container
                let output = tokio::process::Command::new("docker")
                    .args(["rm", &container_id])
                    .output()
                    .await?;

                if !output.status.success() {
                    warn!("Failed to remove container {}: {}",
                          container_id, String::from_utf8_lossy(&output.stderr));
                }

                Ok(())
            }
            Self::TempDirectory { path } => {
                debug!("Removing temporary directory {:?}", path);
                if let Err(e) = std::fs::remove_dir_all(&path) {
                    warn!("Failed to remove temp directory {:?}: {}", path, e);
                }
                Ok(())
            }
        }
    }
}

/// System resource monitor for tracking usage during tests
pub struct ResourceMonitor {
    initial_memory: u64,
    initial_gpu_memory: Vec<u64>,
    start_time: Instant,
}

impl ResourceMonitor {
    pub fn new() -> Result<Self> {
        let mut sys = sysinfo::System::new();
        sys.refresh_memory();
        let initial_memory = sys.used_memory();

        let initial_gpu_memory = Vec::new(); // Would query GPU memory here

        Ok(Self {
            initial_memory,
            initial_gpu_memory,
            start_time: Instant::now(),
        })
    }

    pub fn snapshot(&self) -> ResourceSnapshot {
        let mut sys = sysinfo::System::new();
        sys.refresh_memory();
        let current_memory = sys.used_memory();

        ResourceSnapshot {
            timestamp: Instant::now(),
            memory_used_mb: current_memory / 1024 / 1024,
            memory_delta_mb: ((current_memory as i64) - (self.initial_memory as i64)) / 1024 / 1024,
            gpu_memory_used_mb: Vec::new(), // Would query GPU memory here
            elapsed: self.start_time.elapsed(),
        }
    }
}

#[derive(Debug, Clone)]
pub struct ResourceSnapshot {
    pub timestamp: Instant,
    pub memory_used_mb: u64,
    pub memory_delta_mb: i64,
    pub gpu_memory_used_mb: Vec<u64>,
    pub elapsed: Duration,
}

impl ResourceSnapshot {
    pub fn print_summary(&self) {
        info!("Resource snapshot at {:?}:", self.elapsed);
        info!("  Memory: {} MB (Δ{:+} MB)", self.memory_used_mb, self.memory_delta_mb);
        if !self.gpu_memory_used_mb.is_empty() {
            for (i, gpu_mem) in self.gpu_memory_used_mb.iter().enumerate() {
                info!("  GPU {}: {} MB", i, gpu_mem);
            }
        }
    }
}

/// Network utilities for distributed testing
pub struct NetworkUtils;

impl NetworkUtils {
    /// Find an available port for testing
    pub async fn find_available_port() -> Result<u16> {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
        let port = listener.local_addr()?.port();
        Ok(port)
    }

    /// Check if a service is reachable
    pub async fn wait_for_service(host: &str, port: u16, timeout_secs: u64) -> Result<()> {
        let timeout_duration = Duration::from_secs(timeout_secs);
        let start = Instant::now();

        while start.elapsed() < timeout_duration {
            if let Ok(_) = tokio::net::TcpStream::connect((host, port)).await {
                return Ok(());
            }
            tokio::time::sleep(Duration::from_millis(100)).await;
        }

        anyhow::bail!("Service at {host}:{port} not reachable within {timeout_duration:?}")
    }

    /// Get local IP address for multi-node testing
    pub fn get_local_ip() -> Result<String> {
        let local_ip = local_ip_address::local_ip()?;
        Ok(local_ip.to_string())
    }
}

/// Container management utilities
pub struct ContainerManager {
    runtime: String,
}

impl ContainerManager {
    pub fn new(runtime: String) -> Self {
        Self { runtime }
    }

    /// Start a container and return its ID
    pub async fn start_container(
        &self,
        image: &str,
        ports: &[u16],
        env_vars: &[(&str, &str)],
        volumes: &[(&str, &str)],
    ) -> Result<String> {
        let mut args: Vec<String> = vec!["run".to_string(), "-d".to_string()];

        // Add port mappings
        for port in ports {
            args.push("-p".to_string());
            args.push(format!("{port}:{port}"));
        }

        // Add environment variables
        for (key, value) in env_vars {
            args.push("-e".to_string());
            args.push(format!("{key}={value}"));
        }

        // Add volume mounts
        for (host_path, container_path) in volumes {
            args.push("-v".to_string());
            args.push(format!("{host_path}:{container_path}"));
        }

        args.push(image.to_string());

        let output = tokio::process::Command::new(&self.runtime)
            .args(&args)
            .output()
            .await?;

        if !output.status.success() {
            anyhow::bail!("Failed to start container: {}",
                         String::from_utf8_lossy(&output.stderr));
        }

        let container_id = String::from_utf8(output.stdout)?.trim().to_string();
        info!("Started container {} from image {}", container_id, image);

        Ok(container_id)
    }

    /// Wait for container to be ready
    pub async fn wait_for_container_ready(
        &self,
        container_id: &str,
        timeout_secs: u64,
    ) -> Result<()> {
        let timeout_duration = Duration::from_secs(timeout_secs);
        let start = Instant::now();

        while start.elapsed() < timeout_duration {
            let output = tokio::process::Command::new(&self.runtime)
                .args(["ps", "-q", "--filter", &format!("id={container_id}")])
                .output()
                .await?;

            if output.status.success() && !output.stdout.is_empty() {
                return Ok(());
            }

            tokio::time::sleep(Duration::from_millis(500)).await;
        }

        anyhow::bail!("Container {container_id} not ready within {timeout_duration:?}")
    }

    /// Get container logs
    pub async fn get_container_logs(&self, container_id: &str) -> Result<String> {
        let output = tokio::process::Command::new(&self.runtime)
            .args(["logs", container_id])
            .output()
            .await?;

        Ok(String::from_utf8_lossy(&output.stdout).to_string())
    }

    /// Stop and remove container
    pub async fn cleanup_container(&self, container_id: &str) -> Result<()> {
        // Stop container
        let _ = tokio::process::Command::new(&self.runtime)
            .args(["stop", container_id])
            .output()
            .await;

        // Remove container
        let _ = tokio::process::Command::new(&self.runtime)
            .args(["rm", container_id])
            .output()
            .await;

        Ok(())
    }
}

/// Test data management utilities
pub struct TestDataManager {
    pub base_path: std::path::PathBuf,
}

impl TestDataManager {
    pub fn new(base_path: std::path::PathBuf) -> Self {
        Self { base_path }
    }

    /// Create a temporary directory for test data
    pub async fn create_temp_dir(&self, prefix: &str) -> Result<std::path::PathBuf> {
        let temp_dir = tempfile::tempdir_in(&self.base_path)?;
        let path = temp_dir.into_path();

        // Create subdirectory with prefix
        let test_dir = path.join(format!("{}-{}", prefix, Uuid::new_v4()));
        std::fs::create_dir_all(&test_dir)?;

        Ok(test_dir)
    }

    /// Save test artifacts for debugging
    pub async fn save_artifacts(
        &self,
        test_name: &str,
        artifacts: &[(&str, &[u8])],
    ) -> Result<std::path::PathBuf> {
        let artifacts_dir = self.base_path.join("artifacts").join(test_name);
        std::fs::create_dir_all(&artifacts_dir)?;

        for (filename, data) in artifacts {
            let file_path = artifacts_dir.join(filename);
            std::fs::write(&file_path, data)?;
        }

        Ok(artifacts_dir)
    }

    /// Load test dataset
    pub async fn load_test_dataset(&self, dataset_name: &str) -> Result<Vec<u8>> {
        let dataset_path = self.base_path.join("datasets").join(dataset_name);
        let data = std::fs::read(dataset_path)?;
        Ok(data)
    }
}

/// Performance assertion utilities
pub struct PerformanceAssert;

impl PerformanceAssert {
    /// Assert that operation completes within time limit
    pub async fn assert_duration_max<F, Fut>(
        operation: F,
        max_duration: Duration,
        description: &str,
    ) -> Result<()>
    where
        F: FnOnce() -> Fut,
        Fut: std::future::Future<Output = Result<()>>,
    {
        let start = Instant::now();

        match timeout(max_duration, operation()).await {
            Ok(result) => {
                let elapsed = start.elapsed();
                info!("{} completed in {:?} (limit: {:?})", description, elapsed, max_duration);
                result
            }
            Err(_) => {
                anyhow::bail!("{description} exceeded time limit of {max_duration:?}")
            }
        }
    }

    /// Assert that throughput meets minimum requirement
    pub fn assert_throughput_min(
        operations: usize,
        duration: Duration,
        min_ops_per_sec: f64,
        description: &str,
    ) -> Result<()> {
        let actual_ops_per_sec = operations as f64 / duration.as_secs_f64();

        if actual_ops_per_sec >= min_ops_per_sec {
            info!("{} throughput: {:.2} ops/sec (required: {:.2})",
                  description, actual_ops_per_sec, min_ops_per_sec);
            Ok(())
        } else {
            anyhow::bail!("{description} throughput too low: {actual_ops_per_sec:.2} ops/sec (required: {min_ops_per_sec:.2})")
        }
    }

    /// Assert memory usage stays within limits
    pub fn assert_memory_max(
        current_mb: u64,
        max_mb: u64,
        description: &str,
    ) -> Result<()> {
        if current_mb <= max_mb {
            info!("{} memory usage: {} MB (limit: {} MB)",
                  description, current_mb, max_mb);
            Ok(())
        } else {
            anyhow::bail!("{description} memory usage too high: {current_mb} MB (limit: {max_mb} MB)")
        }
    }
}

/// Retry utilities for flaky operations
pub struct RetryUtils;

impl RetryUtils {
    /// Retry an async operation with exponential backoff
    pub async fn retry_with_backoff<F, Fut, T>(
        operation: F,
        max_attempts: usize,
        initial_delay: Duration,
        max_delay: Duration,
        description: &str,
    ) -> Result<T>
    where
        F: Fn() -> Fut,
        Fut: std::future::Future<Output = Result<T>>,
    {
        let mut attempts = 0;
        let mut delay = initial_delay;

        loop {
            attempts += 1;

            match operation().await {
                Ok(result) => {
                    if attempts > 1 {
                        info!("{} succeeded after {} attempts", description, attempts);
                    }
                    return Ok(result);
                }
                Err(e) if attempts >= max_attempts => {
                    error!("{} failed after {} attempts: {}", description, attempts, e);
                    return Err(e);
                }
                Err(e) => {
                    warn!("{} attempt {} failed: {}", description, attempts, e);
                    tokio::time::sleep(delay).await;
                    delay = std::cmp::min(delay * 2, max_delay);
                }
            }
        }
    }
}

/// Test synchronization utilities
pub struct TestSync;

impl TestSync {
    /// Wait for multiple async operations to complete
    pub async fn wait_all<T>(futures: Vec<impl std::future::Future<Output = Result<T>>>) -> Result<Vec<T>> {
        let results = futures::future::try_join_all(futures).await?;
        Ok(results)
    }

    /// Race multiple operations and return the first successful result
    pub async fn race<T>(futures: Vec<impl std::future::Future<Output = Result<T>> + Unpin>) -> Result<T> {
        let (result, _index, _remaining) = futures::future::select_all(futures).await;
        result
    }
}