test_suite.py

#!/usr/bin/env python3
"""
Test script for math-eval dataset generation pipeline.

This script runs comprehensive tests to ensure all components work correctly.
"""

import os
import sys
import subprocess
import tempfile
import logging

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

def run_command(cmd, description="", check_output=True):
    """Run a command and return success status."""
    logger.info(f"Testing: {description}")
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        if check_output and result.returncode == 0:
            logger.info(f"✓ {description} - PASSED")
            return True
        else:
            logger.error(f"✗ {description} - FAILED")
            logger.error(f"STDERR: {result.stderr}")
            return False
    except subprocess.CalledProcessError as e:
        logger.error(f"✗ {description} - FAILED")
        logger.error(f"Error: {e}")
        logger.error(f"STDERR: {e.stderr}")
        return False

def test_equation_generation():
    """Test basic equation generation."""
    with tempfile.TemporaryDirectory() as temp_dir:
        output_file = os.path.join(temp_dir, "test_equations.txt")
        
        # Test 2-variable equations
        cmd = ["python", "equation_generator.py", "--output_file", output_file, "--num", "5", "--vars", "2"]
        if not run_command(cmd, "Generate 2-variable equations"):
            return False
        
        # Check if file exists and has content
        if not os.path.exists(output_file):
            logger.error("Output file not created")
            return False
        
        with open(output_file, 'r') as f:
            lines = f.readlines()
            if len(lines) != 5:
                logger.error(f"Expected 5 lines, got {len(lines)}")
                return False
        
        # Test 3-variable equations
        output_file_3 = os.path.join(temp_dir, "test_equations_3var.txt")
        cmd = ["python", "equation_generator.py", "--output_file", output_file_3, "--num", "3", "--vars", "3"]
        if not run_command(cmd, "Generate 3-variable equations"):
            return False
        
        return True

def test_verification():
    """Test equation verification."""
    # Test with provided sample files
    if os.path.exists("two-vars.txt"):
        cmd = ["python", "verifier.py", "--file", "two-vars.txt"]
        if not run_command(cmd, "Verify two-vars.txt"):
            return False
    
    if os.path.exists("three-vars.txt"):
        cmd = ["python", "verifier.py", "--file", "three-vars.txt"]  
        if not run_command(cmd, "Verify three-vars.txt"):
            return False
    
    return True

def test_visual_generation():
    """Test visual equation generation."""
    if not os.path.exists("colored_icons_final"):
        logger.warning("Icon directory not found, skipping visual tests")
        return True
    
    with tempfile.TemporaryDirectory() as temp_dir:
        # First generate some equations
        equations_file = os.path.join(temp_dir, "test_equations.txt")
        cmd = ["python", "equation_generator.py", "--output_file", equations_file, "--num", "2", "--vars", "2"]
        if not run_command(cmd, "Generate test equations for visual"):
            return False
        
        # Test character-only generation
        char_dir = os.path.join(temp_dir, "char_only")
        cmd = ["python", "generate_ocr_custom.py", "--equations_file", equations_file, "--output_dir", char_dir]
        if not run_command(cmd, "Generate character-only visuals"):
            return False
        
        # Test icon-only generation
        icon_dir = os.path.join(temp_dir, "icon_only")
        cmd = ["python", "generate_visual_questions.py", "--equations_file", equations_file, "--icon_dir", "colored_icons_final", "--output_dir", icon_dir]
        if not run_command(cmd, "Generate icon-only visuals"):
            return False
        
        # Test counting questions
        counting_dir = os.path.join(temp_dir, "counting")
        cmd = ["python", "generate_counting_questions.py", "--equations_file", equations_file, "--icon_folder", "colored_icons_final", "--output_dir", counting_dir]
        if not run_command(cmd, "Generate counting questions"):
            return False
        
        return True

def test_pipeline():
    """Test the complete pipeline."""
    with tempfile.TemporaryDirectory() as temp_dir:
        cmd = ["python", "run_pipeline.py", "--num_equations", "3", "--num_vars", "2", "--task", "equations", "--output_dir", temp_dir]
        if not run_command(cmd, "Run pipeline (equations only)"):
            return False
        
        # Check if output directory structure is created
        expected_dirs = [
            os.path.join(temp_dir, "equations"),
            os.path.join(temp_dir, "logs")
        ]
        
        for dir_path in expected_dirs:
            if not os.path.exists(dir_path):
                logger.error(f"Expected directory not created: {dir_path}")
                return False
        
        return True

def test_imports():
    """Test that all required modules can be imported."""
    modules_to_test = [
        "equation_generator",
        "verifier", 
        "generate_ocr_custom",
        "generate_visual_questions",
        "generate_counting_questions",
        "generate_partial_visual_questions",
        "run_pipeline"
    ]
    
    for module in modules_to_test:
        try:
            cmd = ["python", "-c", f"import {module}"]
            if not run_command(cmd, f"Import {module}", check_output=True):
                return False
        except Exception as e:
            logger.error(f"Failed to test import of {module}: {e}")
            return False
    
    return True

def test_inference_setup():
    """Test inference system setup."""
    logger.info("Testing inference system...")
    
    # Check if inference config exists
    if not os.path.exists("inference_config.json"):
        logger.error("inference_config.json not found")
        return False
    
    # Check if run_inference.py exists
    if not os.path.exists("run_inference.py"):
        logger.error("run_inference.py not found")
        return False
    
    # Test inference script help
    cmd = ["python", "run_inference.py", "--help"]
    if not run_command(cmd, "Testing run_inference.py help"):
        return False
    
    # Check inference directory structure
    inference_dir = "inference"
    if not os.path.exists(inference_dir):
        logger.error("inference directory not found")
        return False
    
    # Test inference test script
    cmd = ["python", "test_inference.py", "--quick"]
    if not run_command(cmd, "Testing inference system"):
        return False
    
    logger.info("✓ Inference setup test - PASSED")
    return True

def main():
    """Run all tests."""
    logger.info("Starting math-eval test suite...")
    
    tests = [
        ("Import Tests", test_imports),
        ("Equation Generation", test_equation_generation),
        ("Verification", test_verification),
        ("Visual Generation", test_visual_generation),
        ("Pipeline", test_pipeline),
        ("Inference Setup", test_inference_setup)
    ]
    
    results = {}
    for test_name, test_func in tests:
        logger.info(f"\n{'='*50}")
        logger.info(f"Running {test_name}")
        logger.info('='*50)
        
        try:
            results[test_name] = test_func()
        except Exception as e:
            logger.error(f"Test {test_name} failed with exception: {e}")
            results[test_name] = False
    
    # Print summary
    logger.info(f"\n{'='*50}")
    logger.info("TEST SUMMARY")
    logger.info('='*50)
    
    passed = 0
    total = len(results)
    
    for test_name, result in results.items():
        status = "PASSED" if result else "FAILED"
        logger.info(f"{test_name}: {status}")
        if result:
            passed += 1
    
    logger.info(f"\nOverall: {passed}/{total} tests passed")
    
    if passed == total:
        logger.info("🎉 All tests passed! The math-eval system is working correctly.")
        sys.exit(0)
    else:
        logger.error("❌ Some tests failed. Please check the logs above.")
        sys.exit(1)

if __name__ == "__main__":
    main()