run_tests.py

#!/usr/bin/env python3
"""
Test runner for R2BC experimental output unit tests

This script runs all unit tests for the policy checkpoint saving/loading
and gif generation functionality.
"""

import unittest
import sys
from pathlib import Path
import argparse

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

def discover_and_run_tests(test_dir="tests", pattern="test_*.py", verbosity=2):
    """
    Discover and run all tests in the specified directory
    
    Args:
        test_dir: Directory containing test files
        pattern: Pattern to match test files
        verbosity: Verbosity level (0=quiet, 1=normal, 2=verbose)
    
    Returns:
        bool: True if all tests passed, False otherwise
    """
    
    # Create test loader
    loader = unittest.TestLoader()
    
    # Discover tests
    start_dir = Path(test_dir)
    if not start_dir.exists():
        print(f"Error: Test directory '{test_dir}' does not exist!")
        return False
    
    suite = loader.discover(str(start_dir), pattern=pattern)
    
    # Create test runner
    runner = unittest.TextTestRunner(verbosity=verbosity, buffer=True)
    
    # Run tests
    print(f"Running tests from {test_dir} with pattern {pattern}")
    print("=" * 70)
    
    result = runner.run(suite)
    
    # Print summary
    print("\n" + "=" * 70)
    print("TEST SUMMARY")
    print("=" * 70)
    print(f"Tests run: {result.testsRun}")
    print(f"Failures: {len(result.failures)}")
    print(f"Errors: {len(result.errors)}")
    print(f"Skipped: {len(result.skipped) if hasattr(result, 'skipped') else 0}")
    
    if result.failures:
        print(f"\nFAILURES ({len(result.failures)}):")
        for test, traceback in result.failures:
            print(f"  - {test}")
    
    if result.errors:
        print(f"\nERRORS ({len(result.errors)}):")
        for test, traceback in result.errors:
            print(f"  - {test}")
    
    success = len(result.failures) == 0 and len(result.errors) == 0
    
    if success:
        print("\n✅ ALL TESTS PASSED!")
    else:
        print(f"\n❌ {len(result.failures) + len(result.errors)} TEST(S) FAILED!")
    
    return success

def run_specific_test_module(module_name, verbosity=2):
    """
    Run tests from a specific module
    
    Args:
        module_name: Name of the test module (e.g., 'test_policy_checkpoints')
        verbosity: Verbosity level
        
    Returns:
        bool: True if all tests passed, False otherwise
    """
    
    try:
        # Import the test module
        test_module = __import__(f"tests.{module_name}", fromlist=[module_name])
        
        # Create test suite from module
        loader = unittest.TestLoader()
        suite = loader.loadTestsFromModule(test_module)
        
        # Run tests
        runner = unittest.TextTestRunner(verbosity=verbosity, buffer=True)
        result = runner.run(suite)
        
        return len(result.failures) == 0 and len(result.errors) == 0
        
    except ImportError as e:
        print(f"Error: Could not import test module '{module_name}': {e}")
        return False

def main():
    """Main entry point for test runner"""
    
    parser = argparse.ArgumentParser(description="Run R2BC experimental output tests")
    parser.add_argument(
        "--module", "-m", 
        help="Run tests from specific module (e.g., test_policy_checkpoints)"
    )
    parser.add_argument(
        "--verbose", "-v", 
        action="store_true", 
        help="Increase verbosity"
    )
    parser.add_argument(
        "--quiet", "-q", 
        action="store_true", 
        help="Reduce verbosity"
    )
    parser.add_argument(
        "--pattern", "-p", 
        default="test_*.py",
        help="Pattern to match test files (default: test_*.py)"
    )
    
    args = parser.parse_args()
    
    # Determine verbosity level
    verbosity = 2  # Default
    if args.quiet:
        verbosity = 0
    elif args.verbose:
        verbosity = 2
    
    # Run tests
    if args.module:
        success = run_specific_test_module(args.module, verbosity)
    else:
        success = discover_and_run_tests(pattern=args.pattern, verbosity=verbosity)
    
    # Exit with appropriate code
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()