ds_mcp_template/run_tests.py

#!/usr/bin/env python3
"""
Test Runner for MCP Template

This script provides an easy way to run different types of tests
with appropriate configurations and environment setup.
"""

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


def run_command(cmd, cwd=None, env=None):
    """Run a command and return the result."""
    try:
        result = subprocess.run(
            cmd,
            shell=True,
            cwd=cwd or os.getcwd(),
            env=env or os.environ.copy(),
            capture_output=True,
            text=True
        )
        return result.returncode == 0, result.stdout, result.stderr
    except Exception as e:
        return False, "", str(e)


def setup_environment():
    """Set up environment variables for testing."""
    env = os.environ.copy()

    # Load .env file if it exists
    env_file = Path(".env")
    if env_file.exists():
        try:
            from dotenv import load_dotenv
            load_dotenv()
            env.update(os.environ.copy())
        except ImportError:
            print("Warning: python-dotenv not installed, .env file not loaded")

    return env


def run_unit_tests(args):
    """Run unit tests."""
    print("🧪 Running Unit Tests...")

    cmd = ["python", "-m", "pytest", "tests/unit/", "-v"]
    if args.coverage:
        cmd.extend(["--cov=src", "--cov-report=html", "--cov-report=term"])

    success, stdout, stderr = run_command(" ".join(cmd))
    print(stdout)
    if stderr:
        print(stderr)

    return success


def run_integration_tests(args):
    """Run integration tests."""
    print("🔗 Running Integration Tests...")

    # Check for API keys
    has_openai = bool(os.getenv("OPENAI_API_KEY"))
    has_anthropic = bool(os.getenv("ANTHROPIC_API_KEY"))
    has_grok = bool(os.getenv("GROK_API_KEY"))

    if not any([has_openai, has_anthropic, has_grok]):
        print("⚠️  Warning: No API keys found. Integration tests will be skipped.")
        print("   Set OPENAI_API_KEY, ANTHROPIC_API_KEY, or GROK_API_KEY environment variables.")

    cmd = ["python", "-m", "pytest", "tests/integration/", "-v", "-s"]

    if args.coverage:
        cmd.extend(["--cov=src", "--cov-report=html", "--cov-report=term"])

    if args.skip_slow:
        cmd.append("-m")
        cmd.append("not slow")

    success, stdout, stderr = run_command(" ".join(cmd))
    print(stdout)
    if stderr:
        print(stderr)

    return success


def run_all_tests(args):
    """Run all tests."""
    print("🚀 Running All Tests...")

    success1 = run_unit_tests(args)
    success2 = run_integration_tests(args)

    return success1 and success2


def run_specific_test(args):
    """Run a specific test file or test function."""
    print(f"🎯 Running Specific Test: {args.test_path}")

    cmd = ["python", "-m", "pytest", args.test_path, "-v", "-s"]

    if args.coverage:
        cmd.extend(["--cov=src", "--cov-report=html", "--cov-report=term"])

    success, stdout, stderr = run_command(" ".join(cmd))
    print(stdout)
    if stderr:
        print(stderr)

    return success


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Test Runner for MCP Template",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Run all tests
  python run_tests.py all

  # Run unit tests only
  python run_tests.py unit

  # Run integration tests only
  python run_tests.py integration

  # Run specific test file
  python run_tests.py specific tests/integration/test_mcp_integration.py

  # Run with coverage
  python run_tests.py all --coverage

  # Run integration tests but skip slow ones
  python run_tests.py integration --skip-slow
        """
    )

    parser.add_argument(
        "command",
        choices=["unit", "integration", "all", "specific"],
        help="Type of tests to run"
    )

    parser.add_argument(
        "test_path",
        nargs="?",
        help="Path to specific test file or test function (for 'specific' command)"
    )

    parser.add_argument(
        "--coverage",
        action="store_true",
        help="Generate coverage report"
    )

    parser.add_argument(
        "--skip-slow",
        action="store_true",
        help="Skip slow tests (integration tests only)"
    )

    args = parser.parse_args()

    # Set up environment
    env = setup_environment()

    # Change to the correct working directory
    os.chdir(Path(__file__).parent)

    # Run the appropriate test command
    if args.command == "unit":
        success = run_unit_tests(args)
    elif args.command == "integration":
        success = run_integration_tests(args)
    elif args.command == "all":
        success = run_all_tests(args)
    elif args.command == "specific":
        if not args.test_path:
            print("❌ Error: test_path is required for 'specific' command")
            sys.exit(1)
        success = run_specific_test(args)
    else:
        print("❌ Unknown command")
        success = False

    # Exit with appropriate code
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()