testlint_sdk/runtime_coverage/

python.rs

use crate::platform::{process_exists, signal_process, ProcessSignal};
use crate::runtime_coverage::{CoverageSummary, RuntimeCoverageResult};
use std::fs;
use std::path::Path;
use std::process::Command;
use std::thread;
use std::time::Duration;

pub struct PythonRuntimeCoverage;

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

impl PythonRuntimeCoverage {
    pub fn new() -> Self {
        PythonRuntimeCoverage
    }

    /// Attach to a running Python process and collect coverage continuously
    pub fn attach_and_collect(&self, pid: u32) -> Result<RuntimeCoverageResult, String> {
        println!("🐍 Attaching coverage to Python process PID: {}", pid);
        println!("📊 Collecting coverage continuously (press Ctrl+C to stop)...");

        // Check if process exists
        self.check_process_exists(pid)?;

        // Create a Python script that will inject coverage into the target process
        let inject_script = self.create_injection_script(pid)?;

        // Run the injection script
        println!("💉 Injecting coverage.py into process...");
        let output = Command::new("python3")
            .arg(&inject_script)
            .output()
            .map_err(|e| format!("Failed to run injection script: {}", e))?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(format!("Coverage injection failed: {}", stderr));
        }

        println!("✓ Coverage injected successfully");
        println!("⏳ Collecting coverage... Press Ctrl+C to stop and save");

        // Wait for Ctrl+C
        use std::sync::atomic::{AtomicBool, Ordering};
        use std::sync::Arc;

        let running = Arc::new(AtomicBool::new(true));
        let r = running.clone();

        ctrlc::set_handler(move || {
            r.store(false, Ordering::SeqCst);
        })
        .expect("Error setting Ctrl-C handler");

        let start_time = std::time::Instant::now();

        while running.load(Ordering::SeqCst) {
            thread::sleep(Duration::from_millis(100));
        }

        let duration_secs = start_time.elapsed().as_secs();

        // Signal the process to save coverage (send SIGUSR1 on Unix, file trigger on Windows)
        println!("\n📝 Signaling process to save coverage data...");

        #[cfg(unix)]
        {
            signal_process(pid, ProcessSignal::User1)?;
        }

        #[cfg(windows)]
        {
            // On Windows, create a trigger file that the injected code watches for
            let trigger_file = std::env::temp_dir().join(format!("coverage_trigger_{}.txt", pid));
            fs::write(&trigger_file, "save")?;
            println!(
                "ℹ️  Created trigger file (Windows alternative to SIGUSR1): {}",
                trigger_file.display()
            );
        }

        // Wait a bit for coverage to be saved
        thread::sleep(Duration::from_secs(2));

        // Find and parse the coverage file
        let coverage_file = self.find_coverage_file(pid)?;
        let summary = self.parse_coverage_summary(&coverage_file)?;

        // Clean up injection script
        let _ = fs::remove_file(&inject_script);

        Ok(RuntimeCoverageResult {
            language: "Python".to_string(),
            pid,
            duration_secs,
            coverage_file,
            summary,
        })
    }

    /// Check if the process exists
    fn check_process_exists(&self, pid: u32) -> Result<(), String> {
        if !process_exists(pid)? {
            return Err(format!("Process {} not found", pid));
        }
        Ok(())
    }

    /// Create a Python script that injects coverage into the target process
    fn create_injection_script(&self, pid: u32) -> Result<String, String> {
        let temp_dir = std::env::temp_dir();
        let script_path = temp_dir.join(format!("coverage_inject_{}.py", pid));
        let script_path_str = script_path.to_string_lossy().to_string();

        let script_content = format!(
            r#"#!/usr/bin/env python3
"""
Runtime coverage injection script.
Injects coverage.py into a running Python process using gdb.
"""
import os
import sys
import time
import signal
import tempfile

PID = {pid}

# Python code to inject into the target process
COVERAGE_CODE = '''
import coverage
import signal
import tempfile
import atexit

# Start coverage
cov = coverage.Coverage(data_file=".coverage.{{}}".format(os.getpid()), branch=True)
cov.start()

print("[Coverage] Coverage started for PID {{}}".format(os.getpid()), file=sys.stderr)

# Save coverage on SIGUSR1
def save_coverage(signum, frame):
    print("[Coverage] Saving coverage data...", file=sys.stderr)
    cov.stop()
    cov.save()
    print("[Coverage] Coverage saved to .coverage.{{}}".format(os.getpid()), file=sys.stderr)

# Also save on exit
atexit.register(lambda: (cov.stop(), cov.save()))
signal.signal(signal.SIGUSR1, save_coverage)
'''

def inject_via_pyrasite():
    """Try using pyrasite to inject code."""
    try:
        import pyrasite
        print(f"Injecting coverage into PID {{PID}} using pyrasite...")

        # Create temp file with coverage code
        with open(f'/tmp/cov_inject_{{PID}}.py', 'w') as f:
            f.write(COVERAGE_CODE)

        # Inject
        pyrasite.inject(PID, f'/tmp/cov_inject_{{PID}}.py')
        print("✓ Coverage injected successfully")
        return True
    except ImportError:
        print("pyrasite not installed, trying gdb method...")
        return False
    except Exception as e:
        print(f"pyrasite injection failed: {{e}}")
        return False

def inject_via_gdb():
    """Fallback: use gdb to inject Python code."""
    import subprocess

    print(f"Injecting coverage into PID {{PID}} using gdb...")

    gdb_commands = f'''
set pagination off
call (void*)PyGILState_Ensure()
call (int)PyRun_SimpleString("{{COVERAGE_CODE.replace(chr(10), '\\\\n')}}")
call (void)PyGILState_Release($1)
detach
quit
'''

    with open(f'/tmp/gdb_commands_{{PID}}.txt', 'w') as f:
        f.write(gdb_commands)

    try:
        result = subprocess.run(
            ['gdb', '-p', str(PID), '-batch', '-x', f'/tmp/gdb_commands_{{PID}}.txt'],
            capture_output=True,
            text=True,
            timeout=30
        )

        if result.returncode == 0:
            print("✓ Coverage injected via gdb")
            return True
        else:
            print(f"gdb injection failed: {{result.stderr}}")
            return False
    except subprocess.TimeoutExpired:
        print("gdb injection timed out")
        return False
    except FileNotFoundError:
        print("gdb not found. Please install gdb.")
        return False
    finally:
        # Cleanup
        try:
            os.remove(f'/tmp/gdb_commands_{{PID}}.txt')
        except:
            pass

# Try pyrasite first, fall back to gdb
if not inject_via_pyrasite():
    if not inject_via_gdb():
        print("ERROR: Failed to inject coverage. Install pyrasite or gdb:")
        print("  pip install pyrasite")
        print("  or: apt-get install gdb / brew install gdb")
        sys.exit(1)

print("Coverage collection running continuously...")
print("Process will save coverage when receiving SIGUSR1 signal")
"#,
            pid = pid
        );

        fs::write(&script_path, script_content)
            .map_err(|e| format!("Failed to create injection script: {}", e))?;

        // Make executable
        #[cfg(unix)]
        {
            Command::new("chmod")
                .args(["+x", &script_path_str])
                .output()
                .map_err(|e| format!("Failed to make script executable: {}", e))?;
        }

        Ok(script_path_str)
    }

    /// Find the coverage file generated by the process
    fn find_coverage_file(&self, pid: u32) -> Result<String, String> {
        let coverage_file = format!(".coverage.{}", pid);

        if Path::new(&coverage_file).exists() {
            return Ok(coverage_file);
        }

        // Try current directory
        if Path::new(".coverage").exists() {
            return Ok(".coverage".to_string());
        }

        Err(format!(
            "Coverage file not found. Expected .coverage.{} or .coverage",
            pid
        ))
    }

    /// Parse coverage summary from coverage file
    fn parse_coverage_summary(&self, coverage_file: &str) -> Result<CoverageSummary, String> {
        println!("📊 Generating coverage report...");

        // Run coverage report to get summary
        let output = Command::new("coverage")
            .args(["report", "--data-file", coverage_file])
            .output()
            .map_err(|e| format!("Failed to run coverage report: {}", e))?;

        if !output.status.success() {
            return Err(format!(
                "Coverage report failed: {}",
                String::from_utf8_lossy(&output.stderr)
            ));
        }

        let report = String::from_utf8_lossy(&output.stdout);

        // Parse the TOTAL line
        // Format: TOTAL    1234   123   90%
        let total_line = report
            .lines()
            .find(|line| line.starts_with("TOTAL"))
            .ok_or("Could not find TOTAL line in coverage report")?;

        let parts: Vec<&str> = total_line.split_whitespace().collect();
        if parts.len() < 4 {
            return Err("Invalid coverage report format".to_string());
        }

        let total_lines = parts[1]
            .parse::<usize>()
            .map_err(|_| "Invalid total lines")?;
        let missed_lines = parts[2]
            .parse::<usize>()
            .map_err(|_| "Invalid missed lines")?;
        let covered_lines = total_lines - missed_lines;
        let coverage_percentage = parts[3]
            .trim_end_matches('%')
            .parse::<f64>()
            .map_err(|_| "Invalid coverage percentage")?;

        Ok(CoverageSummary {
            total_lines,
            covered_lines,
            coverage_percentage,
            total_branches: None,
            covered_branches: None,
            branch_percentage: None,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_runtime_coverage_new() {
        let coverage = PythonRuntimeCoverage::new();
        assert_eq!(std::mem::size_of_val(&coverage), 0);
    }

    #[test]
    fn test_runtime_coverage_default() {
        let coverage = PythonRuntimeCoverage;
        assert_eq!(std::mem::size_of_val(&coverage), 0);
    }
}