test_agent.py

import sys
import importlib
from pathlib import Path


def check_python_version():
    """Check if Python version is compatible"""
    print("Checking Python version...", end=" ")

    version = sys.version_info
    if version.major < 3 or (version.major == 3 and version.minor < 8):
        print(f" FAIL")
        print(f"   Current: Python {version.major}.{version.minor}.{version.micro}")
        print(f"   Required: Python 3.8+")
        return False

    print(f" OK (Python {version.major}.{version.minor}.{version.micro})")
    return True


def check_dependencies():
    """Check if all required dependencies are installed"""
    print("\nChecking dependencies...")

    required = {
        'numpy': 'numpy',
        'pandas': 'pandas',
        'sklearn': 'scikit-learn',
        'matplotlib': 'matplotlib',
        'joblib': 'joblib',
        'fastapi':'uvicorn[standard]',
        'python-multipart':'python-multipart',
        'websockets': 'websockets',
        'aiofiles':'aiofiles'
    }

    optional = {
        'seaborn': 'seaborn',
        'optuna': 'optuna',
        'xgboost':'xgboost',
        'lightgbm':'lightgbm'
    }

    all_ok = True

    # Check required
    for module, package in required.items():
        try:
            importlib.import_module(module)
            print(f"    {package}")
        except ImportError:
            print(f"   {package} - REQUIRED")
            all_ok = False

    # Check optional
    for module, package in optional.items():
        try:
            importlib.import_module(module)
            print(f" {package}")
        except ImportError:
            print(f" {package} - OPTIONAL (recommended)")

    return all_ok


def check_file_structure():
    """Check if all required files are present"""
    print("\nChecking file structure...", end=" ")

    required_files = [
        'ml_agent.py',
        'config.py',
        'example_usage.py',
        'requirements.txt',
        'README.md'
    ]

    missing = []
    for file in required_files:
        if not Path(file).exists():
            missing.append(file)

    if missing:
        print("FAIL")
        print("   Missing files:")
        for file in missing:
            print(f"      - {file}")
        return False

    print(" OK")
    return True


def test_basic_functionality():
    """Test basic ML Agent functionality"""
    print("\nTesting basic functionality...")

    try:
        from ml_agent import MLAgent
        import pandas as pd
        import numpy as np

        print(" Imports successful")

        # Create tiny test dataset
        np.random.seed(42)
        df = pd.DataFrame({
            'feature1': np.random.randn(50),
            'feature2': np.random.randn(50),
            'target': np.random.randint(0, 2, 50)
        })

        # Save to temp file
        test_dir = Path('test_outputs')
        test_dir.mkdir(exist_ok=True)

        test_file = test_dir / 'test_data.csv'
        df.to_csv(test_file, index=False)

        print(" Test data created")

        # Try to initialize agent
        agent = MLAgent(
            data_path=str(test_file),
            output_dir=str(test_dir / 'outputs'),
            target_col='target',
            problem_type='classification'
        )

        print(" Agent initialized")

        # Test data loading
        agent.load_data()
        print(" Data loading works")

        # Test fingerprinting
        fingerprint = agent.dataset_fingerprint
        if fingerprint and len(fingerprint) == 32:
            print(f" Fingerprinting works ({fingerprint[:8]}...)")
        else:
            print(" Fingerprinting may have issues")

        # Test analysis
        agent.analyze_data()
        print(" Data analysis works")

        # Cleanup
        import shutil
        shutil.rmtree(test_dir)
        print(" Cleanup complete")

        return True

    except Exception as e:
        print(f"FAIL: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


def run_full_test():
    """Run a complete end-to-end test"""
    print("\nRunning full end-to-end test...")
    print("This may take 1-2 minutes...")

    try:
        from ml_agent import MLAgent
        from example_usage import generate_iris_dataset

        # Generate test data
        df = generate_iris_dataset()

        test_dir = Path('test_outputs')
        test_dir.mkdir(exist_ok=True)

        test_file = test_dir / 'test_iris.csv'
        df.to_csv(test_file, index=False)

        print(" Test data generated")

        # Run agent
        agent = MLAgent(
            data_path=str(test_file),
            output_dir=str(test_dir / 'outputs')
        )

        print(" Running agent...")
        agent.run()

        # Verify outputs
        output_dir = test_dir / 'outputs'

        checks = [
            ('Model file', output_dir / 'models' / 'final_model.joblib'),
            ('Feature importance plot', output_dir / 'plots' / 'feature_importance.png'),
            ('Model comparison plot', output_dir / 'plots' / 'metric_comparison.png'),
            ('Overview report', output_dir / 'reports' / 'overview.md'),
            ('Data analysis report', output_dir / 'reports' / 'data_analysis.md'),
            ('Modeling report', output_dir / 'reports' / 'modeling.md'),
            ('Results report', output_dir / 'reports' / 'results.md'),
        ]

        all_ok = True
        print("\n   Verifying outputs:")
        for name, path in checks:
            if path.exists():
                print(f" {name}")
            else:
                print(f" {name} - NOT FOUND")
                all_ok = False

        # Test model loading
        import joblib
        model_pkg = joblib.load(output_dir / 'models' / 'final_model.joblib')

        if 'model' in model_pkg and 'preprocessor' in model_pkg:
            print(f" Model loadable")
        else:
            print(f" Model structure incorrect")
            all_ok = False

        # Cleanup
        import shutil
        shutil.rmtree(test_dir)

        if all_ok:
            print("\n Full test PASSED")
        else:
            print("\n  Full test completed with warnings")

        return all_ok

    except Exception as e:
        print(f"   Full test FAILED: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


def main():
    """Run all tests"""
    print("=" * 60)
    print("ML AGENT - INSTALLATION & FUNCTIONALITY TEST")
    print("=" * 60)

    results = []

    # 1. Python version
    results.append(("Python Version", check_python_version()))

    # 2. Dependencies
    results.append(("Dependencies", check_dependencies()))

    # 3. File structure
    results.append(("File Structure", check_file_structure()))

    # 4. Basic functionality
    results.append(("Basic Functionality", test_basic_functionality()))

    # Summary
    print("\n" + "=" * 60)
    print("TEST SUMMARY")
    print("=" * 60)

    for name, passed in results:
        status = "PASS" if passed else "FAIL"
        print(f"{name:.<40} {status}")

    # Overall result
    all_passed = all(result[1] for result in results)

    print("=" * 60)

    if all_passed:
        print("ALL TESTS PASSED")
        print("\nYour ML Agent is ready to use!")
        print("\nNext steps:")
        print("1. Generate sample data: python example_usage.py --generate-data")
        print("2. Run on sample data: python ml_agent.py sample_data/iris.csv")
        print("3. Check the results in outputs/")

        # Ask if user wants full test
        print("\n" + "=" * 60)
        response = input("\nRun full end-to-end test? (y/n): ").strip().lower()

        if response == 'y':
            success = run_full_test()

            print("\n" + "=" * 60)
            if success:
                print("COMPLETE SUCCESS - Everything is working perfectly!")
            else:
                print("  Some issues found - Check the output above")
            print("=" * 60)
    else:
        print("SOME TESTS FAILED")
        print("\nPlease fix the issues above before using the ML Agent.")
        print("\nCommon fixes:")
        print("- Install missing dependencies: pip install -r requirements.txt")
        print("- Ensure all files are in the same directory")
        print("- Check Python version is 3.8+")

    print()


if __name__ == "__main__":
    main()