test_tool_functionality.py

#!/usr/bin/env python3
"""
Test tool functionality in current environment.

This script tests if the benchmarking tools can run with minimal synthetic data
to validate the configuration system and tool implementations.
"""

import pandas as pd
import anndata as ad
import numpy as np
from pathlib import Path
import tempfile
import sys

# Add the project root to the path
sys.path.insert(0, str(Path(__file__).parent))

from benchmarking.tools.scanpy_tool import create_scanpy_tool
from benchmarking.tools.meld_tool import create_meld_tool  
from benchmarking.tools.seurat_tool import create_seurat_tool


def create_synthetic_data():
    """Create minimal synthetic AnnData for testing."""
    np.random.seed(42)
    
    # Create synthetic expression data
    n_cells = 200
    n_genes = 100
    
    # Create raw count data (integer counts)
    X_raw = np.random.negative_binomial(5, 0.5, size=(n_cells, n_genes)).astype(int)
    
    # Create normalized data (log-transformed)
    import scanpy as sc
    
    # Create cell metadata
    obs = pd.DataFrame({
        'cell_id': [f'cell_{i}' for i in range(n_cells)],
        'condition': ['Young'] * 100 + ['Old'] * 100,
        'sample': [f'sample_{i//50}' for i in range(n_cells)]
    })
    obs.index = obs['cell_id']
    
    # Create gene metadata
    var = pd.DataFrame({
        'gene_id': [f'gene_{i}' for i in range(n_genes)],
        'gene_name': [f'Gene_{i}' for i in range(n_genes)]
    })
    var.index = var['gene_id']
    
    # Create AnnData object with raw counts
    adata = ad.AnnData(X=X_raw.astype(float), obs=obs, var=var)
    
    # Store raw counts and create normalized data
    adata.raw = adata  # Store raw counts
    
    # Normalize and log-transform the data
    sc.pp.normalize_total(adata, target_sum=1e4)
    sc.pp.log1p(adata)
    
    # Store raw counts in a layer too
    adata.layers['raw_counts'] = X_raw.astype(float)
    
    # Add PCA for tools that need manifold coordinates (like MELD)
    sc.pp.pca(adata, n_comps=50)
    
    return adata


def test_scanpy_tool():
    """Test the scanpy tool with synthetic data."""
    print("=== Testing scanpy tool ===")
    
    # Create synthetic data
    adata = create_synthetic_data()
    
    # Create configuration
    config = {
        'conditions': ['Young', 'Old'],
        'grouping_column': 'condition',
        'sample_column': 'sample',
        'raw_counts_layer': 'raw_counts'  # Use raw counts layer for proper scanpy analysis
    }
    
    # Create tool and test
    tool = create_scanpy_tool()
    
    with tempfile.TemporaryDirectory() as temp_dir:
        output_dir = Path(temp_dir)
        
        try:
            obs_results, var_results = tool.run_analysis(adata, config, output_dir)
            
            if obs_results is not None:
                print(f"✓ obs_results: {obs_results.shape}")
                print(f"  Columns: {list(obs_results.columns)}")
            else:
                print("❌ No obs_results returned")
                
            if var_results is not None:
                print(f"✓ var_results: {var_results.shape}")
                print(f"  Columns: {list(var_results.columns)}")
            else:
                print("❌ No var_results returned")
                
            if obs_results is not None or var_results is not None:
                print("✓ Scanpy tool working!")
                return True
            else:
                print("❌ Scanpy tool failed")
                return False
                
        except Exception as e:
            print(f"❌ Scanpy tool error: {e}")
            return False


def test_meld_tool():
    """Test the MELD tool with synthetic data."""
    print("=== Testing MELD tool ===")
    
    try:
        # Create synthetic data
        adata = create_synthetic_data()
        
        # Create configuration
        config = {
            'conditions': ['Young', 'Old'],
            'grouping_column': 'condition',
            'sample_column': 'sample',
            'raw_counts_layer': 'raw_counts'
        }
        
        # Create tool and test  
        import logging
        logging.basicConfig(level=logging.INFO)
        tool = create_meld_tool()
        
        with tempfile.TemporaryDirectory() as temp_dir:
            output_dir = Path(temp_dir)
            
            obs_results, var_results = tool.run_analysis(adata, config, output_dir)
            
            print(f"MELD tool.run_analysis returned: obs_results={type(obs_results)}, var_results={type(var_results)}")
            
            if obs_results is not None:
                print(f"✓ obs_results: {obs_results.shape}")
                print(f"  Columns: {list(obs_results.columns)}")
                # Check if it has MELD-specific columns
                meld_columns = [col for col in obs_results.columns if 'meld' in col.lower()]
                if meld_columns:
                    print(f"  ✓ MELD columns found: {meld_columns}")
                    print("✓ MELD tool working!")
                    return True
                else:
                    print(f"  ❌ No MELD-specific columns found")
                    print("❌ MELD tool failed - no results")
                    return False
            else:
                print("❌ No obs_results returned")
                print("❌ MELD tool failed - no output")
                return False
                
            if var_results is not None:
                print(f"✓ var_results: {var_results.shape}")
                print(f"  Columns: {list(var_results.columns)}")
            else:
                print("ℹ️ No var_results (expected for MELD DA)")
                
    except Exception as e:
        print(f"❌ MELD tool error: {e}")
        import traceback
        print(f"Full traceback:\n{traceback.format_exc()}")
        return False


def test_seurat_tool():
    """Test the Seurat tool with synthetic data."""
    print("=== Testing Seurat tool ===")
    
    try:
        # Create synthetic data
        adata = create_synthetic_data()
        
        # Create configuration
        config = {
            'conditions': ['Young', 'Old'],
            'grouping_column': 'condition',
            'sample_column': 'sample',
            'raw_counts_layer': 'raw_counts'
        }
        
        # Create tool and test
        tool = create_seurat_tool()
        
        with tempfile.TemporaryDirectory() as temp_dir:
            output_dir = Path(temp_dir)
            
            obs_results, var_results = tool.run_analysis(adata, config, output_dir)
            
            if obs_results is not None:
                print(f"✓ obs_results: {obs_results.shape}")
                print(f"  Columns: {list(obs_results.columns)}")
            else:
                print("ℹ️ No obs_results (expected for Seurat DE)")
                
            if var_results is not None:
                print(f"✓ var_results: {var_results.shape}")
                print(f"  Columns: {list(var_results.columns)}")
            else:
                print("❌ No var_results returned")
                
            if var_results is not None:
                print("✓ Seurat tool working!")
                return True
            else:
                print("❌ Seurat tool failed")
                return False
                
    except Exception as e:
        print(f"❌ Seurat tool error: {e}")
        import traceback
        print(f"Traceback: {traceback.format_exc()}")
        return False


def test_current_environment():
    """Test if current environment has necessary packages."""
    print("=== Testing current environment ===")
    
    packages = [
        'pandas', 'numpy', 'scanpy', 'anndata', 'yaml', 'pathlib'
    ]
    
    missing = []
    for pkg in packages:
        try:
            __import__(pkg)
            print(f"✓ {pkg} available")
        except ImportError:
            print(f"❌ {pkg} missing")
            missing.append(pkg)
    
    if not missing:
        print("✓ All required packages available in current environment")
        return True
    else:
        print(f"❌ Missing packages: {missing}")
        return False


def main():
    """Run all tests."""
    print("Testing tool functionality across environments...\n")
    
    # Test environment
    env_ok = test_current_environment()
    print()
    
    if not env_ok:
        print("❌ Current environment missing dependencies. Cannot proceed.")
        return False
    
    # Test tools
    print("=== Testing tools ===")
    tools_tested = 0
    tools_working = 0
    
    # Test scanpy (should work in current environment)
    scanpy_ok = test_scanpy_tool()
    tools_tested += 1
    if scanpy_ok:
        tools_working += 1
    print()
    
    # Test MELD (should work in current environment)
    meld_ok = test_meld_tool()
    tools_tested += 1
    if meld_ok:
        tools_working += 1
    print()
    
    # Test Seurat (requires renv environment)
    seurat_ok = test_seurat_tool()
    tools_tested += 1
    if seurat_ok:
        tools_working += 1
    print()
    
    # Summary
    print("=== Tool Test Summary ===")
    print(f"✓ Environment setup: {'PASS' if env_ok else 'FAIL'}")
    print(f"✓ Scanpy tool: {'PASS' if scanpy_ok else 'FAIL'}")
    print(f"✓ MELD tool: {'PASS' if meld_ok else 'FAIL'}")
    print(f"✓ Seurat tool: {'PASS' if seurat_ok else 'FAIL'}")
    print(f"\nTools working: {tools_working}/{tools_tested}")
    
    if tools_working > 0:
        print(f"\n✅ Tool functionality confirmed! {tools_working} out of {tools_tested} tools working.")
        if tools_working < tools_tested:
            print("Some tools may require additional environment setup or dependencies.")
        return True
    else:
        print("\n❌ No tools working. Check environment setup and dependencies.")
        return False


if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)