test_inference.py

#!/usr/bin/env python3
"""
Test script for inference functionality.
Validates that all inference components work correctly.
"""

import os
import sys
import json
import subprocess
import argparse
import logging
from pathlib import Path

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def run_command(cmd, description, capture_output=True):
    """Run a command and return success status."""
    logger.info(f"Testing: {description}")
    try:
        if capture_output:
            result = subprocess.run(cmd, shell=True, check=True, capture_output=True, text=True)
            return True, result.stdout
        else:
            result = subprocess.run(cmd, shell=True, check=True)
            return True, ""
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed: {description}")
        if capture_output and e.stderr:
            logger.error(f"Error: {e.stderr}")
        return False, ""

def test_config_file():
    """Test if inference config exists and is valid."""
    config_path = "inference_config.json"
    if not os.path.exists(config_path):
        logger.error(f"Config file not found: {config_path}")
        return False
    
    try:
        with open(config_path, 'r') as f:
            config = json.load(f)
        
        required_sections = ['api_models', 'open_source', 'datasets', 'output_dir']
        for section in required_sections:
            if section not in config:
                logger.error(f"Missing config section: {section}")
                return False
        
        logger.info("✅ Config file is valid")
        return True
    except json.JSONDecodeError as e:
        logger.error(f"Invalid JSON in config: {e}")
        return False

def test_inference_scripts():
    """Test that all inference scripts exist and have basic syntax."""
    inference_dir = Path("inference")
    if not inference_dir.exists():
        logger.error("Inference directory not found")
        return False
    
    # Find all Python files
    python_files = list(inference_dir.rglob("*.py"))
    
    if not python_files:
        logger.error("No inference scripts found")
        return False
    
    logger.info(f"Found {len(python_files)} inference scripts")
    
    # Test syntax of each script
    failed = 0
    for script in python_files:
        cmd = f"python -m py_compile {script}"
        success, _ = run_command(cmd, f"Syntax check: {script.name}")
        if not success:
            failed += 1
    
    if failed == 0:
        logger.info("✅ All inference scripts have valid syntax")
        return True
    else:
        logger.error(f"❌ {failed} scripts have syntax errors")
        return False

def test_run_inference_script():
    """Test the main run_inference.py script."""
    if not os.path.exists("run_inference.py"):
        logger.error("run_inference.py not found")
        return False
    
    # Test help output
    success, output = run_command("python run_inference.py --help", "run_inference.py help")
    if not success:
        return False
    
    # Check that help contains expected options
    required_args = ['--task', '--model_type', '--inference_type']
    for arg in required_args:
        if arg not in output:
            logger.error(f"Missing argument in help: {arg}")
            return False
    
    logger.info("✅ run_inference.py script is functional")
    return True

def test_sample_datasets():
    """Test if sample datasets exist for inference."""
    datasets = [
        "three-vars/icon_only/metadata.csv",
        "three-vars/char_only/metadata.csv", 
        "three-vars/icon_partial/metadata.csv"
    ]
    
    missing = []
    for dataset in datasets:
        if not os.path.exists(dataset):
            missing.append(dataset)
    
    if missing:
        logger.warning(f"Missing sample datasets: {missing}")
        logger.warning("Run data generation first: make run-small")
        return False
    else:
        logger.info("✅ Sample datasets available")
        return True

def test_dry_run_inference():
    """Test inference with dry run mode."""
    if not test_sample_datasets():
        logger.info("Skipping dry run test - no sample data")
        return True
    
    # Test dry run with different configurations
    test_configs = [
        {
            "task": "visual_equation_solving",
            "dataset": "icon_only",
            "model_type": "api",
            "api_model": "openai",
            "inference_type": "direct",
            "api_key": "dummy-key"
        }
    ]
    
    for config in test_configs:
        cmd_parts = ["python run_inference.py", "--dry_run"]
        for key, value in config.items():
            cmd_parts.append(f"--{key} {value}")
        
        cmd = " ".join(cmd_parts)
        success, output = run_command(cmd, f"Dry run: {config['task']}/{config.get('dataset', '')}")
        
        if not success:
            return False
        
        if "Would execute:" not in output:
            logger.error("Dry run did not produce expected output")
            return False
    
    logger.info("✅ Dry run inference tests passed")
    return True

def main():
    parser = argparse.ArgumentParser(description="Test inference functionality")
    parser.add_argument('--quick', action='store_true', help='Run only quick tests')
    args = parser.parse_args()
    
    logger.info("🧪 Testing Math-Eval Inference System")
    logger.info("=" * 50)
    
    tests = [
        ("Config File", test_config_file),
        ("Inference Scripts", test_inference_scripts),
        ("Run Inference Script", test_run_inference_script),
        ("Sample Datasets", test_sample_datasets),
    ]
    
    if not args.quick:
        tests.append(("Dry Run Inference", test_dry_run_inference))
    
    passed = 0
    total = len(tests)
    
    for test_name, test_func in tests:
        logger.info(f"\n📋 Testing: {test_name}")
        if test_func():
            passed += 1
        else:
            logger.error(f"❌ Test failed: {test_name}")
    
    logger.info(f"\n📊 Test Results: {passed}/{total} passed")
    
    if passed == total:
        logger.info("🎉 All inference tests passed!")
        logger.info("\nNext steps:")
        logger.info("1. Configure API keys in inference_config.json")
        logger.info("2. Generate datasets: make run-small")
        logger.info("3. Run inference: python run_inference.py --help")
        return True
    else:
        logger.error("❌ Some tests failed. Please fix issues before proceeding.")
        return False

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