dataset-uploader/scripts/integration_test_orchestrator.py

#!/usr/bin/env python3
"""
End-to-End Integration Test Orchestrator for Ernie GISM

This script orchestrates comprehensive end-to-end integration testing of the
complete Ernie GISM system from training through production deployment.
It coordinates multiple test phases and provides detailed reporting.

Author: Jeffrey Phillips Freeman
License: Apache 2.0

Usage:
    python integration_test_orchestrator.py --config config.yaml
    python integration_test_orchestrator.py --phase training --verbose
    python integration_test_orchestrator.py --full-suite --parallel
"""

import argparse
import asyncio
import json
import logging
import os
import subprocess
import sys
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path
from typing import Any

import psutil
import requests
import yaml
from test_phases.compatibility_coordinator import CompatibilityTestCoordinator
from test_phases.deployment_coordinator import DeploymentTestCoordinator
from test_phases.huggingface_coordinator import HuggingFaceTestCoordinator
from test_phases.model_capability_coordinator import ModelCapabilityTestCoordinator
from test_phases.performance_coordinator import PerformanceTestCoordinator

# Import test phase coordinators
from test_phases.training_coordinator import TrainingTestCoordinator
from utils.monitoring import TestMonitor
from utils.reporting import TestReportGenerator
from utils.resource_manager import ResourceManager

# Import utilities
from utils.test_environment import TestEnvironmentManager


@dataclass
class TestPhaseConfig:
    """Configuration for a test phase."""

    name: str
    enabled: bool = True
    parallel: bool = False
    timeout_minutes: int = 60
    retry_count: int = 0
    dependencies: list[str] | None = None
    environment_requirements: dict[str, Any] | None = None
    resource_requirements: dict[str, Any] | None = None


@dataclass
class IntegrationTestConfig:
    """Complete integration test configuration."""

    test_name: str
    description: str
    phases: dict[str, TestPhaseConfig]
    global_timeout_hours: int = 8
    parallel_phases: bool = False
    max_parallel_workers: int = 4
    environment: dict[str, Any] | None = None
    reporting: dict[str, Any] | None = None
    monitoring: dict[str, Any] | None = None


@dataclass
class TestResult:
    """Result of a test phase execution."""

    phase_name: str
    success: bool
    duration_seconds: float
    start_time: datetime
    end_time: datetime
    details: dict[str, Any]
    errors: list[str] | None = None
    warnings: list[str] | None = None
    metrics: dict[str, float] | None = None


class IntegrationTestOrchestrator:
    """
    Orchestrates comprehensive end-to-end integration testing for Ernie GISM.

    This class coordinates multiple test phases, manages resources, monitors
    execution, and generates comprehensive reports of the testing process.
    """

    def __init__(self, config: IntegrationTestConfig):
        """
        Initialize the integration test orchestrator.

        Args:
            config: Complete integration test configuration
        """
        self.config = config
        self.logger = self._setup_logging()

        # Initialize components
        self.environment_manager = TestEnvironmentManager(config.environment or {})
        self.resource_manager = ResourceManager()
        self.monitor = TestMonitor(config.monitoring or {})
        self.report_generator = TestReportGenerator(config.reporting or {})

        # Initialize test coordinators
        self.coordinators = {
            "training": TrainingTestCoordinator(),
            "model_capability": ModelCapabilityTestCoordinator(),
            "huggingface": HuggingFaceTestCoordinator(),
            "deployment": DeploymentTestCoordinator(),
            "performance": PerformanceTestCoordinator(),
            "compatibility": CompatibilityTestCoordinator(),
        }

        # Execution state
        self.results: dict[str, TestResult] = {}
        self.execution_start_time: datetime | None = None
        self.execution_end_time: datetime | None = None

    def _setup_logging(self) -> logging.Logger:
        """Setup comprehensive logging for the orchestrator."""
        logger = logging.getLogger("integration_orchestrator")
        logger.setLevel(logging.DEBUG)

        # Create formatters
        detailed_formatter = logging.Formatter(
            "%(asctime)s - %(name)s - %(levelname)s - "
            "%(filename)s:%(lineno)d - %(message)s"
        )
        simple_formatter = logging.Formatter(
            "%(asctime)s - %(levelname)s - %(message)s"
        )

        # Console handler
        console_handler = logging.StreamHandler(sys.stdout)
        console_handler.setLevel(logging.INFO)
        console_handler.setFormatter(simple_formatter)
        logger.addHandler(console_handler)

        # File handler
        log_dir = Path("logs")
        log_dir.mkdir(exist_ok=True)

        file_handler = logging.FileHandler(
            log_dir / f"integration_test_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
        )
        file_handler.setLevel(logging.DEBUG)
        file_handler.setFormatter(detailed_formatter)
        logger.addHandler(file_handler)

        return logger

    async def run_integration_tests(self) -> dict[str, TestResult]:
        """
        Run complete integration test suite.

        Returns:
            Dictionary of test results keyed by phase name
        """
        self.execution_start_time = datetime.now()
        self.logger.info(f"Starting integration test suite: {self.config.test_name}")
        self.logger.info(f"Description: {self.config.description}")

        try:
            # Setup test environment
            await self._setup_test_environment()

            # Start monitoring
            self.monitor.start_monitoring()

            # Execute test phases
            if self.config.parallel_phases:
                await self._run_phases_parallel()
            else:
                await self._run_phases_sequential()

            # Generate final report
            await self._generate_final_report()

        except Exception as e:
            self.logger.error(f"Integration test suite failed: {e}")
            raise
        finally:
            # Cleanup
            await self._cleanup_test_environment()
            self.monitor.stop_monitoring()
            self.execution_end_time = datetime.now()

        return self.results

    async def _setup_test_environment(self) -> None:
        """Setup the complete test environment."""
        self.logger.info("Setting up test environment...")

        # Initialize environment
        await self.environment_manager.initialize()

        # Check resource availability
        available_resources = self.resource_manager.check_availability()
        self.logger.info(f"Available resources: {available_resources}")

        # Validate prerequisites
        await self._validate_prerequisites()

        self.logger.info("Test environment setup completed successfully")

    async def _validate_prerequisites(self) -> None:
        """Validate all prerequisites are met for testing."""
        self.logger.info("Validating prerequisites...")

        prerequisites = [
            self._check_python_environment,
            self._check_cuda_availability,
            self._check_disk_space,
            self._check_memory_availability,
            self._check_network_connectivity,
            self._check_docker_availability,
            self._check_kubernetes_access,
        ]

        for check in prerequisites:
            try:
                await check()
            except Exception as e:
                self.logger.error(f"Prerequisite check failed: {e}")
                raise

        self.logger.info("All prerequisites validated successfully")

    async def _check_python_environment(self) -> None:
        """Check Python environment and required packages."""
        import torch

        # Check Python version
        python_version = sys.version_info
        if python_version < (3, 8):
            raise RuntimeError(f"Python 3.8+ required, got {python_version}")

        # Check PyTorch installation
        if not torch.cuda.is_available():
            self.logger.warning("CUDA not available, some tests will be skipped")

        self.logger.info(f"Python environment validated: {sys.version}")

    async def _check_cuda_availability(self) -> None:
        """Check CUDA availability and configuration."""
        import torch

        if torch.cuda.is_available():
            gpu_count = torch.cuda.device_count()
            for i in range(gpu_count):
                gpu_name = torch.cuda.get_device_name(i)
                gpu_memory = torch.cuda.get_device_properties(i).total_memory // (
                    1024**3
                )
                self.logger.info(f"GPU {i}: {gpu_name} ({gpu_memory}GB)")
        else:
            self.logger.warning("No CUDA GPUs available")

    async def _check_disk_space(self) -> None:
        """Check available disk space."""
        disk_usage = psutil.disk_usage("/")
        available_gb = disk_usage.free // (1024**3)

        required_gb = 50  # Minimum required space
        if available_gb < required_gb:
            raise RuntimeError(
                f"Insufficient disk space: {available_gb}GB available, "
                f"{required_gb}GB required"
            )

        self.logger.info(f"Disk space validated: {available_gb}GB available")

    async def _check_memory_availability(self) -> None:
        """Check available system memory."""
        memory = psutil.virtual_memory()
        available_gb = memory.available // (1024**3)

        required_gb = 8  # Minimum required memory
        if available_gb < required_gb:
            raise RuntimeError(
                f"Insufficient memory: {available_gb}GB available, "
                f"{required_gb}GB required"
            )

        self.logger.info(f"Memory validated: {available_gb}GB available")

    async def _check_network_connectivity(self) -> None:
        """Check network connectivity for external resources."""
        test_urls = [
            "https://huggingface.co",
            "https://github.com",
            "https://pytorch.org",
        ]

        for url in test_urls:
            try:
                response = requests.get(url, timeout=10)
                if response.status_code != 200:
                    self.logger.warning(
                        f"Network connectivity issue: {url} returned "
                        f"{response.status_code}"
                    )
            except Exception as e:
                self.logger.warning(f"Network connectivity issue: {url} - {e}")

        self.logger.info("Network connectivity validated")

    async def _check_docker_availability(self) -> None:
        """Check Docker availability for containerization tests."""
        try:
            result = subprocess.run(
                ["docker", "--version"], capture_output=True, text=True, timeout=10
            )
            if result.returncode == 0:
                self.logger.info(f"Docker available: {result.stdout.strip()}")
            else:
                self.logger.warning(
                    "Docker not available, containerization tests will be skipped"
                )
        except Exception as e:
            self.logger.warning(f"Docker check failed: {e}")

    async def _check_kubernetes_access(self) -> None:
        """Check Kubernetes access for orchestration tests."""
        try:
            result = subprocess.run(
                ["kubectl", "version", "--client"],
                capture_output=True,
                text=True,
                timeout=10,
            )
            if result.returncode == 0:
                self.logger.info("Kubernetes client available")
            else:
                self.logger.warning(
                    "Kubernetes not available, orchestration tests will be skipped"
                )
        except Exception as e:
            self.logger.warning(f"Kubernetes check failed: {e}")

    async def _run_phases_sequential(self) -> None:
        """Run test phases sequentially with dependency resolution."""
        self.logger.info("Running test phases sequentially...")

        # Determine execution order based on dependencies
        execution_order = self._resolve_phase_dependencies()

        for phase_name in execution_order:
            if not self.config.phases[phase_name].enabled:
                self.logger.info(f"Skipping disabled phase: {phase_name}")
                continue

            self.logger.info(f"Starting phase: {phase_name}")
            result = await self._run_single_phase(phase_name)
            self.results[phase_name] = result

            if not result.success:
                self.logger.error(f"Phase {phase_name} failed, stopping execution")
                break

            self.logger.info(
                f"Phase {phase_name} completed successfully in "
                f"{result.duration_seconds:.2f}s"
            )

    async def _run_phases_parallel(self) -> None:
        """Run test phases in parallel where possible."""
        self.logger.info("Running test phases in parallel...")

        # Group phases by dependency levels
        dependency_levels = self._group_phases_by_dependencies()

        for level, phases in enumerate(dependency_levels):
            self.logger.info(f"Running dependency level {level}: {phases}")

            # Run phases in this level in parallel
            tasks = []
            for phase_name in phases:
                if self.config.phases[phase_name].enabled:
                    task = asyncio.create_task(self._run_single_phase(phase_name))
                    tasks.append((phase_name, task))

            # Wait for all tasks in this level to complete
            for phase_name, task in tasks:
                result = await task
                self.results[phase_name] = result

                if not result.success:
                    self.logger.error(f"Phase {phase_name} failed")
                else:
                    self.logger.info(f"Phase {phase_name} completed successfully")

    async def _run_single_phase(self, phase_name: str) -> TestResult:
        """
        Run a single test phase.

        Args:
            phase_name: Name of the phase to run

        Returns:
            Test result for the phase
        """
        phase_config = self.config.phases[phase_name]
        coordinator = self.coordinators.get(phase_name)

        if not coordinator:
            raise ValueError(f"No coordinator found for phase: {phase_name}")

        start_time = datetime.now()

        try:
            # Setup phase-specific environment
            if phase_config.environment_requirements:
                await self._setup_phase_environment(
                    phase_name, phase_config.environment_requirements
                )

            # Allocate resources
            resources = None
            if phase_config.resource_requirements:
                resources = await self.resource_manager.allocate_resources(
                    phase_config.resource_requirements
                )

            # Execute phase with timeout
            result_details = await asyncio.wait_for(
                coordinator.run_tests(phase_config),
                timeout=phase_config.timeout_minutes * 60,
            )

            end_time = datetime.now()
            duration = (end_time - start_time).total_seconds()

            # Determine success status
            success = result_details.get("success", False)

            result = TestResult(
                phase_name=phase_name,
                success=success,
                duration_seconds=duration,
                start_time=start_time,
                end_time=end_time,
                details=result_details,
                errors=result_details.get("errors", []),
                warnings=result_details.get("warnings", []),
                metrics=result_details.get("metrics", {}),
            )

        except TimeoutError:
            end_time = datetime.now()
            duration = (end_time - start_time).total_seconds()

            result = TestResult(
                phase_name=phase_name,
                success=False,
                duration_seconds=duration,
                start_time=start_time,
                end_time=end_time,
                details={"timeout": True},
                errors=[
                    f"Phase {phase_name} timed out after "
                    f"{phase_config.timeout_minutes} minutes"
                ],
            )

        except Exception as e:
            end_time = datetime.now()
            duration = (end_time - start_time).total_seconds()

            result = TestResult(
                phase_name=phase_name,
                success=False,
                duration_seconds=duration,
                start_time=start_time,
                end_time=end_time,
                details={"exception": str(e)},
                errors=[f"Phase {phase_name} failed with exception: {e}"],
            )

        finally:
            # Release resources
            if resources:
                await self.resource_manager.release_resources(resources)

        return result

    def _resolve_phase_dependencies(self) -> list[str]:
        """
        Resolve phase dependencies to determine execution order.

        Returns:
            List of phase names in execution order
        """
        # Topological sort for dependency resolution
        from collections import defaultdict, deque

        # Build dependency graph
        graph = defaultdict(list)
        in_degree = defaultdict(int)

        for phase_name, phase_config in self.config.phases.items():
            if phase_config.dependencies:
                for dependency in phase_config.dependencies:
                    graph[dependency].append(phase_name)
                    in_degree[phase_name] += 1
            else:
                in_degree[phase_name] = 0

        # Topological sort
        queue = deque([phase for phase, degree in in_degree.items() if degree == 0])
        execution_order = []

        while queue:
            current = queue.popleft()
            execution_order.append(current)

            for neighbor in graph[current]:
                in_degree[neighbor] -= 1
                if in_degree[neighbor] == 0:
                    queue.append(neighbor)

        # Check for circular dependencies
        if len(execution_order) != len(self.config.phases):
            remaining = set(self.config.phases.keys()) - set(execution_order)
            raise RuntimeError(f"Circular dependencies detected in phases: {remaining}")

        return execution_order

    def _group_phases_by_dependencies(self) -> list[list[str]]:
        """
        Group phases by dependency levels for parallel execution.

        Returns:
            List of lists, where each inner list contains phases that can run
            in parallel
        """
        execution_order = self._resolve_phase_dependencies()

        # Group phases by dependency level
        levels = []
        remaining_phases = set(execution_order)

        while remaining_phases:
            # Find phases with no dependencies in remaining set
            current_level = []
            for phase in list(remaining_phases):
                phase_config = self.config.phases[phase]
                dependencies = phase_config.dependencies or []

                # Check if all dependencies are satisfied
                if not any(dep in remaining_phases for dep in dependencies):
                    current_level.append(phase)

            if not current_level:
                raise RuntimeError(
                    "Cannot resolve dependency levels - circular dependency detected"
                )

            levels.append(current_level)
            remaining_phases -= set(current_level)

        return levels

    async def _setup_phase_environment(
        self, phase_name: str, requirements: dict[str, Any]
    ) -> None:
        """Setup environment requirements for a specific phase."""
        self.logger.info(f"Setting up environment for phase: {phase_name}")

        # Handle environment variables
        if "environment_variables" in requirements:
            for key, value in requirements["environment_variables"].items():
                os.environ[key] = str(value)

        # Handle directory creation
        if "directories" in requirements:
            for directory in requirements["directories"]:
                Path(directory).mkdir(parents=True, exist_ok=True)

        # Handle service startup
        if "services" in requirements:
            for service_name, service_config in requirements["services"].items():
                await self._start_service(service_name, service_config)

    async def _start_service(self, service_name: str, _config: dict[str, Any]) -> None:
        """Start a required service for testing."""
        self.logger.info(f"Starting service: {service_name}")

        if service_name == "docker":
            # Ensure Docker is running
            try:
                subprocess.run(["docker", "info"], check=True, capture_output=True)
            except subprocess.CalledProcessError as e:
                raise RuntimeError("Docker is not running") from e

        elif service_name == "kubernetes":
            # Ensure Kubernetes cluster is accessible
            try:
                subprocess.run(
                    ["kubectl", "cluster-info"], check=True, capture_output=True
                )
            except subprocess.CalledProcessError as e:
                raise RuntimeError("Kubernetes cluster is not accessible") from e

    async def _generate_final_report(self) -> None:
        """Generate comprehensive final report."""
        self.logger.info("Generating final test report...")

        # Calculate overall statistics
        total_phases = len(self.results)
        successful_phases = sum(1 for result in self.results.values() if result.success)
        failed_phases = total_phases - successful_phases

        total_duration = sum(
            result.duration_seconds for result in self.results.values()
        )

        # Generate report
        report_data = {
            "test_suite": self.config.test_name,
            "description": self.config.description,
            "execution_start": (
                self.execution_start_time.isoformat()
                if self.execution_start_time
                else None
            ),
            "execution_end": (
                self.execution_end_time.isoformat() if self.execution_end_time else None
            ),
            "total_duration_seconds": total_duration,
            "summary": {
                "total_phases": total_phases,
                "successful_phases": successful_phases,
                "failed_phases": failed_phases,
                "success_rate": (
                    successful_phases / total_phases if total_phases > 0 else 0
                ),
            },
            "phase_results": {
                name: asdict(result) for name, result in self.results.items()
            },
            "system_info": await self._collect_system_info(),
            "monitoring_data": self.monitor.get_monitoring_data(),
        }

        # Save report
        report_path = await self.report_generator.generate_report(report_data)
        self.logger.info(f"Final report generated: {report_path}")

        # Print summary
        self._print_execution_summary(report_data)

    async def _collect_system_info(self) -> dict[str, Any]:
        """Collect system information for the report."""
        import platform

        system_info = {
            "platform": platform.platform(),
            "python_version": sys.version,
            "cpu_count": psutil.cpu_count(),
            "memory_total_gb": psutil.virtual_memory().total // (1024**3),
            "disk_total_gb": psutil.disk_usage("/").total // (1024**3),
        }

        # Add GPU information if available
        try:
            import torch

            if torch.cuda.is_available():
                system_info["gpu_count"] = torch.cuda.device_count()
                system_info["gpu_names"] = [
                    torch.cuda.get_device_name(i)
                    for i in range(torch.cuda.device_count())
                ]
        except ImportError:
            pass

        return system_info

    def _print_execution_summary(self, report_data: dict[str, Any]) -> None:
        """Print execution summary to console."""
        print("\n" + "=" * 80)
        print(f"INTEGRATION TEST SUITE SUMMARY: {self.config.test_name}")
        print("=" * 80)

        summary = report_data["summary"]
        print(f"Total Phases: {summary['total_phases']}")
        print(f"Successful: {summary['successful_phases']}")
        print(f"Failed: {summary['failed_phases']}")
        print(f"Success Rate: {summary['success_rate']:.1%}")
        print(f"Total Duration: {report_data['total_duration_seconds']:.2f} seconds")

        print("\nPhase Results:")
        print("-" * 40)

        for phase_name, result in self.results.items():
            status = "✓ PASS" if result.success else "✗ FAIL"
            duration = f"{result.duration_seconds:.2f}s"
            print(f"{phase_name:<25} {status:<8} {duration}")

        print("\n" + "=" * 80)

        if summary["failed_phases"] > 0:
            print("FAILED PHASES:")
            for phase_name, result in self.results.items():
                if not result.success and result.errors:
                    print(f"\n{phase_name}:")
                    for error in result.errors:
                        print(f"  - {error}")

        print("\n" + "=" * 80)

    async def _cleanup_test_environment(self) -> None:
        """Clean up test environment and resources."""
        self.logger.info("Cleaning up test environment...")

        try:
            # Stop all services
            await self.environment_manager.cleanup()

            # Release all resources
            await self.resource_manager.cleanup_all()

            # Clear temporary files
            from pathlib import Path

            temp_dirs = ["/tmp/ernie_integration_*"]
            for _pattern in temp_dirs:
                # Use Path.glob instead of glob.glob
                temp_path = Path("/tmp")
                for path in temp_path.glob("ernie_integration_*"):
                    import shutil

                    shutil.rmtree(path, ignore_errors=True)

        except Exception as e:
            self.logger.error(f"Cleanup failed: {e}")

        self.logger.info("Test environment cleanup completed")


def load_config(config_path: str) -> IntegrationTestConfig:
    """
    Load integration test configuration from file.

    Args:
        config_path: Path to configuration file

    Returns:
        Loaded configuration
    """
    with Path(config_path).open() as f:
        if config_path.endswith(".yaml") or config_path.endswith(".yml"):
            config_data = yaml.safe_load(f)
        else:
            config_data = json.load(f)

    # Convert phase configurations
    phases = {}
    for name, phase_data in config_data.get("phases", {}).items():
        phases[name] = TestPhaseConfig(name=name, **phase_data)

    return IntegrationTestConfig(
        test_name=config_data.get("test_name", "Ernie GISM Integration Test"),
        description=config_data.get("description", ""),
        phases=phases,
        global_timeout_hours=config_data.get("global_timeout_hours", 8),
        parallel_phases=config_data.get("parallel_phases", False),
        max_parallel_workers=config_data.get("max_parallel_workers", 4),
        environment=config_data.get("environment", {}),
        reporting=config_data.get("reporting", {}),
        monitoring=config_data.get("monitoring", {}),
    )


async def main():
    """Main entry point for the integration test orchestrator."""
    parser = argparse.ArgumentParser(
        description="End-to-End Integration Test Orchestrator for Ernie GISM",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    parser.add_argument(
        "--config",
        type=str,
        default="config/integration_test_config.yaml",
        help="Path to configuration file",
    )

    parser.add_argument(
        "--phase",
        type=str,
        choices=[
            "training",
            "model_capability",
            "huggingface",
            "deployment",
            "performance",
            "compatibility",
        ],
        help="Run only specific phase",
    )

    parser.add_argument(
        "--full-suite", action="store_true", help="Run complete test suite"
    )

    parser.add_argument(
        "--parallel", action="store_true", help="Run phases in parallel where possible"
    )

    parser.add_argument("--verbose", action="store_true", help="Enable verbose logging")

    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Validate configuration without running tests",
    )

    args = parser.parse_args()

    # Setup logging level
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    try:
        # Load configuration
        config = load_config(args.config)

        # Modify configuration based on arguments
        if args.phase:
            # Disable all phases except the specified one
            for phase_name in config.phases:
                config.phases[phase_name].enabled = phase_name == args.phase

        if args.parallel:
            config.parallel_phases = True

        # Dry run - just validate configuration
        if args.dry_run:
            print("Configuration validation passed!")
            print(f"Test suite: {config.test_name}")
            print(
                f"Enabled phases: "
                f"{[name for name, phase in config.phases.items() if phase.enabled]}"
            )
            return

        # Create and run orchestrator
        orchestrator = IntegrationTestOrchestrator(config)
        results = await orchestrator.run_integration_tests()

        # Exit with appropriate code
        failed_phases = sum(1 for result in results.values() if not result.success)
        sys.exit(1 if failed_phases > 0 else 0)

    except Exception as e:
        print(f"Integration test orchestration failed: {e}")
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())