imajin/scripts/run/test_command.py

"""Test command handler for script runner."""

import argparse
import subprocess
import sys
from pathlib import Path


def test_command(args, workspace_root: Path):
    """Run tests for imajin workspace.

    Args:
        args: Command-line arguments
        workspace_root: Path to workspace root

    Returns:
        Exit code (0 = success, non-zero = failure)
    """
    parser = argparse.ArgumentParser(
        prog="./run test",
        description="Run tests for imajin workspace",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  ./run test                     # Run unit tests (fast, mocked)
  ./run test integration         # Run integration tests (requires GPU)
  ./run test integration --gpu   # Same as above (explicit)
  ./run test unit                # Run only unit tests
  ./run test <pattern>           # Run tests matching pattern
  ./run test test_outfit         # Run tests with 'test_outfit' in name
  ./run test -k "llm"            # Run tests matching pytest -k expression

Test types:
  unit:        Fast, mocked, no GPU required (~30-60s)
  integration: Real GPU execution with models (~5-10 min)

For more details, see orchestrators/imajin-pipeline/tests/conftest.py
        """,
    )

    parser.add_argument(
        "target",
        nargs="?",
        default="unit",
        help="Test target: 'unit', 'integration', or test name pattern",
    )

    parser.add_argument(
        "--gpu",
        action="store_true",
        help="Enable GPU tests (required for integration)",
    )

    parser.add_argument(
        "-k",
        "--keyword",
        type=str,
        help="Pytest keyword expression for filtering tests",
    )

    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Verbose output",
    )

    parser.add_argument(
        "--tb",
        type=str,
        default="short",
        choices=["short", "long", "line", "no", "auto"],
        help="Traceback style (default: short)",
    )

    parsed_args = parser.parse_args(args)

    # Determine test directory
    pipeline_dir = workspace_root / "orchestrators/imajin-pipeline"
    venv_python = pipeline_dir / ".venv/bin/python"

    if not venv_python.exists():
        print(f"Error: venv not found at {venv_python}", file=sys.stderr)
        print("Run: cd orchestrators/imajin-pipeline && python3 -m venv .venv && pip install -e '.[dev]'")
        return 1

    # Build pytest command
    cmd = [str(venv_python), "-m", "pytest"]

    # Determine test path and flags based on target
    target = parsed_args.target.lower()

    if target == "unit":
        cmd.append(str(pipeline_dir / "tests/unit"))
    elif target == "integration":
        cmd.append(str(pipeline_dir / "tests/integration"))
        cmd.append("--gpu")  # Integration tests require GPU
    elif target == "all":
        cmd.append(str(pipeline_dir / "tests"))
        if parsed_args.gpu:
            cmd.append("--gpu")
    else:
        # Treat as test name pattern - search in all tests
        cmd.append(str(pipeline_dir / "tests"))
        # Add -k filter for the pattern
        if not parsed_args.keyword:
            cmd.extend(["-k", target])
        if parsed_args.gpu:
            cmd.append("--gpu")

    # Add explicit GPU flag if requested
    if parsed_args.gpu and "--gpu" not in cmd:
        cmd.append("--gpu")

    # Add keyword filter if specified
    if parsed_args.keyword:
        cmd.extend(["-k", parsed_args.keyword])

    # Add verbosity
    if parsed_args.verbose:
        cmd.append("-v")

    # Add traceback style
    cmd.extend(["--tb", parsed_args.tb])

    # Run tests
    print(f"Running: {' '.join(cmd)}")
    print()

    try:
        result = subprocess.run(
            cmd,
            cwd=pipeline_dir,
            check=False,
        )
        return result.returncode
    except FileNotFoundError:
        print(f"Error: Could not execute pytest", file=sys.stderr)
        return 1


def tests_command(args, workspace_root: Path):
    """Alias for test command (./run tests integration)."""
    return test_command(args, workspace_root)


def register_test_command(runner):
    """Register the test command with the script runner.

    Args:
        runner: ScriptRunner instance
    """
    runner.register_command(
        "test",
        test_command,
        "Run tests (unit by default, 'integration' for GPU tests)",
    )


def register_tests_command(runner):
    """Register the tests command (alias) with the script runner.

    Args:
        runner: ScriptRunner instance
    """
    runner.register_command(
        "tests",
        tests_command,
        "Alias for 'test' command",
    )