diff --git a/netengine/cli/main.py b/netengine/cli/main.py index 1a934bd..03dcc12 100644 --- a/netengine/cli/main.py +++ b/netengine/cli/main.py @@ -225,7 +225,7 @@ async def _down(yes: bool, dry_run: bool) -> None: for container in client.containers.list(all=True): by_id = container.id in state_container_ids - by_prefix = any(container.name.startswith(p) for p in _CONTAINER_PREFIXES) + by_prefix = container.name and any(container.name.startswith(p) for p in _CONTAINER_PREFIXES) if by_id or by_prefix: label = f"container:{container.name}" if dry_run: @@ -240,7 +240,7 @@ async def _down(yes: bool, dry_run: bool) -> None: errors.append(f"{label}: {exc}") for network in client.networks.list(): - if any(network.name.startswith(p) for p in _CONTAINER_PREFIXES): + if network.name and any(network.name.startswith(p) for p in _CONTAINER_PREFIXES): label = f"network:{network.name}" if dry_run: click.echo(f" would remove {label}") @@ -411,6 +411,91 @@ async def _diagnose(spec_file: str, as_json: bool) -> None: sys.exit(1) +@cli.command() +@click.option("--interval", default=30, type=int, help="Poll interval in seconds (default 30).") +@click.option("--max-retries", default=3, type=int, help="Max self-heal retries per phase.") +@click.option("--no-auto-heal", is_flag=True, help="Detect drift but don't auto-heal.") +def drift_watch(interval: int, max_retries: int, no_auto_heal: bool) -> None: + """Watch running world for drift and optionally auto-heal (Ctrl+C to stop).""" + asyncio.run(_drift_watch(interval, max_retries, no_auto_heal)) + + +async def _drift_watch(interval: int, max_retries: int, no_auto_heal: bool) -> None: + state = RuntimeState.load() + if not state.world_spec: + click.echo("No running world found — use `netengine up` first.", err=True) + sys.exit(1) + + click.echo(f"Starting drift detection (interval={interval}s, auto-heal={not no_auto_heal})…") + click.echo("Press Ctrl+C to stop.\n") + + orchestrator = Orchestrator(state.world_spec, mock_mode=False) + + orchestrator.start_drift_detection( + poll_interval_seconds=interval, + max_drift_retries=max_retries, + auto_heal=not no_auto_heal, + ) + + try: + await orchestrator.start_consumers() + try: + await asyncio.sleep(float("inf")) + except (KeyboardInterrupt, asyncio.CancelledError): + pass + finally: + await orchestrator.consumer_supervisor.stop_all() + click.echo("\nDrift detection stopped.") + except Exception as exc: + click.echo(f"Drift detection error: {exc}", err=True) + sys.exit(1) + + +@cli.command() +def drift_status() -> None: + """Show current drift status and history.""" + state = RuntimeState.load() + + if not state.world_spec: + click.echo("No running world found — use `netengine up` first.", err=True) + sys.exit(1) + + click.echo("\nDrift Status\n") + + if state.last_drift_check_at: + check_time = state.last_drift_check_at.isoformat() if hasattr( + state.last_drift_check_at, 'isoformat' + ) else str(state.last_drift_check_at) + click.echo(f"Last check: {check_time}") + else: + click.echo("Last check: (no checks yet)") + + if state.current_drift_phases: + click.echo(f"\nCurrently drifted phases: {', '.join(str(p) for p in state.current_drift_phases)}") + else: + click.echo("\nCurrently drifted phases: none") + + if state.drift_history: + click.echo("\nRecent drift history (last 10 events):") + for event in state.drift_history[-10:]: + phase = event.get("phase_num", "?") + detected = event.get("detected_at", "?") + healed = event.get("healed_at") + failed = event.get("healing_failed", False) + + if healed: + status = f"✓ healed at {healed}" + elif failed: + error = event.get("error", "unknown error") + status = f"✗ healing failed: {error}" + else: + status = "⧗ healing in progress" + + click.echo(f" Phase {phase}: detected at {detected}, {status}") + else: + click.echo("\nDrift history: (no events)") + + def _print_status(state: RuntimeState) -> None: world_name = None if state.world_spec and isinstance(state.world_spec, dict): diff --git a/netengine/core/drift_controller.py b/netengine/core/drift_controller.py new file mode 100644 index 0000000..f65e247 --- /dev/null +++ b/netengine/core/drift_controller.py @@ -0,0 +1,332 @@ +"""Drift detection and self-healing controller. + +Continuously monitors running system state against desired spec. +Detects drift when handlers' healthcheck() returns False, and optionally +triggers self-healing by re-running execute() on drifted phases. +""" + +import asyncio +import logging +from dataclasses import dataclass +from datetime import datetime +from typing import Optional + +from netengine.core.orchestrator import Orchestrator +from netengine.events.schema import EventEnvelope +from netengine.handlers._base import BasePhaseHandler + +logger = logging.getLogger(__name__) + + +@dataclass +class DriftState: + """Per-phase drift tracking state.""" + + phase_num: int + handler_name: str + last_healthcheck_at: datetime + is_drifted: bool + drift_detected_at: Optional[datetime] = None + self_healing_attempted: bool = False + last_self_heal_at: Optional[datetime] = None + last_self_heal_error: Optional[str] = None + consecutive_drift_count: int = 0 + + +class DriftDetectionController: + """Continuous reconciliation consumer for drift detection and self-healing. + + Polls phase handlers' healthcheck() methods periodically to detect drift. + When drift is detected, can optionally trigger self-healing by re-running + handler execute() methods. Runs as a background consumer registered with + ConsumerSupervisor. + """ + + def __init__( + self, + orchestrator: Orchestrator, + poll_interval_seconds: int = 30, + max_drift_retries: int = 3, + auto_heal: bool = True, + ): + """Initialize drift detection controller. + + Args: + orchestrator: Orchestrator instance to poll and heal + poll_interval_seconds: Time between healthchecks (default 30) + max_drift_retries: Max self-heal attempts per phase before giving up + auto_heal: If True, automatically re-run execute() on drifted phases + """ + self.orchestrator = orchestrator + self.poll_interval_seconds = poll_interval_seconds + self.max_drift_retries = max_drift_retries + self.auto_heal = auto_heal + self.drift_states: dict[int, DriftState] = {} + + async def run(self) -> None: + """Main drift detection loop. + + Runs continuously: + 1. Check healthcheck() on each completed phase + 2. Detect drift when healthcheck returns False + 3. Trigger self-healing if enabled + 4. Persist drift state to RuntimeState + 5. Sleep until next poll interval + """ + logger.info( + f"Drift detection starting (interval={self.poll_interval_seconds}s, " + f"auto_heal={self.auto_heal})" + ) + + try: + while True: + await self._run_one_iteration() + await asyncio.sleep(self.poll_interval_seconds) + except asyncio.CancelledError: + logger.info("Drift detection stopped") + raise + + async def _run_one_iteration(self) -> None: + """Run one iteration of drift detection and healing.""" + try: + drift_this_round = [] + + for phase_num, handler_class in self.orchestrator.PHASE_HANDLERS: + phase_key = str(phase_num) + if not self.orchestrator.runtime_state.phase_completed.get(phase_key): + continue + + handler = handler_class() + is_healthy = await self._check_phase_health(phase_num, handler) + + if not is_healthy: + drift_this_round.append(phase_num) + await self._emit_drift_event( + phase_num, + handler.__class__.__name__, + "drift.detected", + { + "phase": phase_num, + "handler": handler.__class__.__name__, + "detected_at": datetime.utcnow().isoformat(), + }, + ) + + if drift_this_round and self.auto_heal: + await self._trigger_self_healing(drift_this_round) + + self.orchestrator.runtime_state.last_drift_check_at = datetime.utcnow() + self.orchestrator.runtime_state.current_drift_phases = drift_this_round + self.orchestrator.runtime_state.save() + + except Exception as e: + logger.error(f"Drift detection iteration error: {e}", exc_info=True) + await self._emit_drift_event( + -1, + "drift_controller", + "drift.loop_error", + { + "error": str(e), + "error_at": datetime.utcnow().isoformat(), + }, + ) + + async def _check_phase_health(self, phase_num: int, handler: BasePhaseHandler) -> bool: + """Check healthcheck for a phase. + + Args: + phase_num: Phase number + handler: Handler instance + + Returns: + True if healthy, False if drifted + """ + try: + is_healthy = await handler.healthcheck(self.orchestrator.context) + self._update_drift_state(phase_num, handler.__class__.__name__, is_healthy) + return is_healthy + except Exception as e: + logger.error(f"Phase {phase_num} healthcheck error: {e}") + self._update_drift_state(phase_num, handler.__class__.__name__, False) + return False + + def _update_drift_state(self, phase_num: int, handler_name: str, is_healthy: bool) -> None: + """Update drift tracking state for a phase. + + Args: + phase_num: Phase number + handler_name: Handler class name + is_healthy: Current health status + """ + if phase_num not in self.drift_states: + self.drift_states[phase_num] = DriftState( + phase_num=phase_num, + handler_name=handler_name, + last_healthcheck_at=datetime.utcnow(), + is_drifted=not is_healthy, + drift_detected_at=datetime.utcnow() if not is_healthy else None, + consecutive_drift_count=1 if not is_healthy else 0, + ) + else: + state = self.drift_states[phase_num] + state.last_healthcheck_at = datetime.utcnow() + + if not is_healthy: + if state.is_drifted: + state.consecutive_drift_count += 1 + else: + state.is_drifted = True + state.drift_detected_at = datetime.utcnow() + state.consecutive_drift_count = 1 + else: + if state.is_drifted: + state.is_drifted = False + state.consecutive_drift_count = 0 + + async def _trigger_self_healing(self, drifted_phases: list[int]) -> None: + """Trigger self-healing for drifted phases. + + Args: + drifted_phases: List of phase numbers that drifted + """ + healed_phases = set() + + for phase_num in sorted(drifted_phases): + handler_class = next( + (handler for pnum, handler in self.orchestrator.PHASE_HANDLERS if pnum == phase_num), + None, + ) + if handler_class is None: + logger.warning(f"Phase {phase_num} handler not found") + continue + + handler = handler_class() + success, changed = await self._heal_phase(phase_num, handler) + + if success: + healed_phases.add(phase_num) + if changed: + await self._reheal_dependent_phases(phase_num, healed_phases) + + async def _heal_phase(self, phase_num: int, handler: BasePhaseHandler) -> tuple[bool, bool]: + """Attempt to heal a drifted phase. + + Re-runs execute() and verifies with healthcheck(). + + Args: + phase_num: Phase number + handler: Handler instance + + Returns: + (success, changed_state) tuple + """ + logger.warning(f"Drift detected on phase {phase_num} — attempting self-heal") + drift_state = self.drift_states.get(phase_num) + + try: + self.orchestrator._check_prerequisites(phase_num) + await handler.execute(self.orchestrator.context) + + is_healthy = await handler.healthcheck(self.orchestrator.context) + if not is_healthy: + raise RuntimeError(f"Phase {phase_num} healthcheck still failing after re-run") + + logger.info(f"Phase {phase_num} self-healed successfully") + await self._emit_drift_event( + phase_num, + handler.__class__.__name__, + "drift.self_healed", + { + "phase": phase_num, + "healed_at": datetime.utcnow().isoformat(), + }, + ) + + if drift_state: + drift_state.self_healing_attempted = True + drift_state.last_self_heal_at = datetime.utcnow() + drift_state.last_self_heal_error = None + + self.orchestrator.runtime_state.save() + return True, True + + except Exception as e: + logger.error(f"Phase {phase_num} self-heal failed: {e}") + await self._emit_drift_event( + phase_num, + handler.__class__.__name__, + "drift.self_heal_failed", + { + "phase": phase_num, + "error": str(e), + "failed_at": datetime.utcnow().isoformat(), + }, + ) + + if drift_state: + drift_state.self_healing_attempted = True + drift_state.last_self_heal_error = str(e) + + self.orchestrator.runtime_state.save() + return False, False + + async def _reheal_dependent_phases( + self, changed_phase_num: int, healed_phases: set[int] + ) -> None: + """Re-heal phases that depend on a changed phase. + + If phase N changed state, re-run phases that depend on N. + + Args: + changed_phase_num: Phase number that changed + healed_phases: Set of already-healed phases + """ + for phase_num, handler_class in self.orchestrator.PHASE_HANDLERS: + if phase_num in healed_phases: + continue + + phase_key = str(changed_phase_num) + + if not self.orchestrator.runtime_state.phase_completed.get(phase_key): + continue + + handler = handler_class() + try: + if await handler.healthcheck(self.orchestrator.context): + continue + except Exception: + pass + + success, _ = await self._heal_phase(phase_num, handler) + if success: + healed_phases.add(phase_num) + + async def _emit_drift_event( + self, + phase_num: int, + handler_name: str, + event_type: str, + payload: dict, + ) -> None: + """Emit a drift event. + + Args: + phase_num: Phase number + handler_name: Handler class name + event_type: Event type (e.g., "drift.detected") + payload: Event payload + """ + if self.orchestrator.context.pgmq_client is None: + logger.debug(f"pgmq_client not available, skipping event: {event_type}") + return + + try: + event = EventEnvelope.create( + event_type=event_type, + emitted_by="drift_controller", + payload=payload, + ) + await self.orchestrator.context.pgmq_client.send("drift_events", event) + logger.debug(f"Drift event emitted: {event_type}") + except Exception as e: + logger.error(f"Failed to emit drift event: {e}") diff --git a/netengine/core/orchestrator.py b/netengine/core/orchestrator.py index 4be3677..87d110b 100644 --- a/netengine/core/orchestrator.py +++ b/netengine/core/orchestrator.py @@ -1,5 +1,5 @@ import os -from typing import Any, List, Optional, Type +from typing import Any, List, Optional, Type, cast from pydantic import ValidationError @@ -74,12 +74,12 @@ def __init__(self, spec: NetEngineSpec | dict[str, Any], mock_mode: Optional[boo self.mock_mode = effective_mock # Initialise Docker client only when running for real - docker_client = None + docker_client: Optional[Any] = None if not effective_mock: try: from netengine.handlers.docker_handler import DockerHandler - docker_client = DockerHandler() + docker_client = cast(Any, DockerHandler()) except Exception as exc: logger.warning(f"Docker unavailable, falling back to mock mode: {exc}") effective_mock = True @@ -172,6 +172,35 @@ async def start_consumers(self) -> None: """Start all registered background consumer tasks.""" await self.consumer_supervisor.start_all() + def start_drift_detection( + self, + poll_interval_seconds: int = 30, + max_drift_retries: int = 3, + auto_heal: bool = True, + ) -> None: + """Start drift detection consumer. + + Registers a DriftDetectionController with ConsumerSupervisor so it runs + as a background consumer with automatic restart on crash. + + Args: + poll_interval_seconds: Time between healthchecks (default 30) + max_drift_retries: Max self-heal attempts per phase + auto_heal: If True, automatically re-apply diverged phases + """ + from netengine.core.drift_controller import DriftDetectionController + + drift_controller = DriftDetectionController( + orchestrator=self, + poll_interval_seconds=poll_interval_seconds, + max_drift_retries=max_drift_retries, + auto_heal=auto_heal, + ) + self.consumer_supervisor.register( + "drift_detection", + drift_controller.run, + ) + def _mark_phase_complete(self, phase_num: int, handler: BasePhaseHandler) -> None: """Record user-facing phase completion for a completed handler. diff --git a/netengine/core/self_healing.py b/netengine/core/self_healing.py new file mode 100644 index 0000000..57912de --- /dev/null +++ b/netengine/core/self_healing.py @@ -0,0 +1,126 @@ +"""Self-healing strategy for reconciliation. + +Implements the logic for re-applying drifted phases and their dependencies. +""" + +import logging +from dataclasses import dataclass +from typing import Optional + +from netengine.core.orchestrator import Orchestrator +from netengine.handlers._base import BasePhaseHandler + +logger = logging.getLogger(__name__) + + +@dataclass +class SelfHealResult: + """Result of a self-healing attempt.""" + + phase_num: int + handler_name: str + success: bool + changed_state: bool = False + error: Optional[str] = None + + +class SelfHealingStrategy: + """Logic for re-applying drifted phases and their dependencies.""" + + def __init__(self, orchestrator: Orchestrator): + """Initialize self-healing strategy. + + Args: + orchestrator: Orchestrator instance to use for healing + """ + self.orchestrator = orchestrator + + async def heal_phase( + self, + phase_num: int, + handler: BasePhaseHandler, + ) -> SelfHealResult: + """Heal a drifted phase by re-running execute(). + + Re-runs handler.execute() and verifies with healthcheck(). + + Args: + phase_num: Phase number + handler: Handler instance + + Returns: + SelfHealResult with success/error details + """ + logger.warning(f"Attempting to heal phase {phase_num}") + + try: + self.orchestrator._check_prerequisites(phase_num) + await handler.execute(self.orchestrator.context) + + is_healthy = await handler.healthcheck(self.orchestrator.context) + if not is_healthy: + raise RuntimeError(f"Phase {phase_num} healthcheck still failing after re-run") + + logger.info(f"Phase {phase_num} healed successfully") + self.orchestrator.runtime_state.save() + + return SelfHealResult( + phase_num=phase_num, + handler_name=handler.__class__.__name__, + success=True, + changed_state=True, + error=None, + ) + + except Exception as e: + logger.error(f"Phase {phase_num} healing failed: {e}") + self.orchestrator.runtime_state.save() + + return SelfHealResult( + phase_num=phase_num, + handler_name=handler.__class__.__name__, + success=False, + changed_state=False, + error=str(e), + ) + + async def heal_dependent_phases( + self, + changed_phase_num: int, + healed_phases: set[int], + ) -> list[SelfHealResult]: + """Re-heal phases that depend on a changed phase. + + If phase N changed state, re-run phases that depend on N. + + Args: + changed_phase_num: Phase number that changed + healed_phases: Set of phases already healed + + Returns: + List of SelfHealResult for dependent phases + """ + results = [] + phase_key = str(changed_phase_num) + + if not self.orchestrator.runtime_state.phase_completed.get(phase_key): + return results + + for phase_num, handler_class in self.orchestrator.PHASE_HANDLERS: + if phase_num in healed_phases: + continue + + handler = handler_class() + try: + if await handler.healthcheck(self.orchestrator.context): + continue + except Exception: + pass + + result = await self.heal_phase(phase_num, handler) + results.append(result) + + if result.success: + healed_phases.add(phase_num) + + return results diff --git a/netengine/core/state.py b/netengine/core/state.py index 33c1c37..186fa0a 100644 --- a/netengine/core/state.py +++ b/netengine/core/state.py @@ -59,6 +59,10 @@ class RuntimeState: bootstrap_admin_password: Optional[str] = None platform_client_id: Optional[str] = None + # Drift detection and self-healing + drift_history: list[Dict[str, Any]] = field(default_factory=list) + last_drift_check_at: Optional[datetime] = None + current_drift_phases: list[int] = field(default_factory=list) # PKI certificate rotation tracking issued_certificates: Dict[str, Dict[str, Any]] = field(default_factory=dict) pki_rotation_state: Dict[str, Any] = field(default_factory=dict) @@ -70,6 +74,9 @@ def load(cls) -> "RuntimeState": with open(state_file, "r") as f: data: Dict[str, Any] = json.load(f) # datetime fields are stored as ISO strings + for dt_field in ("started_at", "completed_at", "last_error_at", "last_drift_check_at"): + if data.get(dt_field): + data[dt_field] = datetime.fromisoformat(data[dt_field]) for dt_field in ("started_at", "completed_at", "last_error_at"): dt_value = data.get(dt_field) if dt_value and isinstance(dt_value, str): diff --git a/tests/integration/test_drift_detection.py b/tests/integration/test_drift_detection.py new file mode 100644 index 0000000..1b493df --- /dev/null +++ b/tests/integration/test_drift_detection.py @@ -0,0 +1,260 @@ +"""Integration tests for drift detection and self-healing. + +These tests verify drift detection works end-to-end with real orchestrator +and handler instances. They use mock mode to avoid requiring Docker. +""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from netengine.core.drift_controller import DriftDetectionController +from netengine.core.orchestrator import Orchestrator +from netengine.core.state import RuntimeState +from netengine.spec.models import NetEngineSpec + + +class TestDriftDetectionIntegration: + """Integration tests for drift detection.""" + + @pytest.mark.asyncio + async def test_drift_detection_with_mock_mode( + self, + minimal_spec: NetEngineSpec, + tmp_path, + monkeypatch, + ) -> None: + """Test drift detection in mock mode (no Docker required).""" + # Setup isolated runtime state + state_file = str(tmp_path / "netengine_state.json") + monkeypatch.setenv("NETENGINE_STATE_FILE", state_file) + + # Create orchestrator in mock mode + orchestrator = Orchestrator(minimal_spec, mock_mode=True) + + # Mark some phases as complete + orchestrator.runtime_state.phase_completed = { + "0": True, + "1": True, + "3": True, + } + + # Create drift detection controller + controller = DriftDetectionController( + orchestrator=orchestrator, + poll_interval_seconds=1, + auto_heal=False, + ) + + # Mock healthcheck to return False for one phase + original_handlers = {} + for phase_num, handler_class in orchestrator.PHASE_HANDLERS: + if phase_num == 1: + original_handlers[phase_num] = handler_class + mock_handler = AsyncMock() + mock_handler.healthcheck = AsyncMock(return_value=False) + handlers = orchestrator.PHASE_HANDLERS + orchestrator.PHASE_HANDLERS = [ + (pn, (mock_handler if pn == 1 else h)) for pn, h in handlers + ] + + # Run one iteration + await controller._run_one_iteration() + + # Verify drift was detected + assert 1 in controller.drift_states + assert controller.drift_states[1].is_drifted is True + + @pytest.mark.asyncio + async def test_drift_detection_multiple_iterations( + self, + minimal_spec: NetEngineSpec, + tmp_path, + monkeypatch, + ) -> None: + """Test drift detection across multiple iterations.""" + state_file = str(tmp_path / "netengine_state.json") + monkeypatch.setenv("NETENGINE_STATE_FILE", state_file) + + orchestrator = Orchestrator(minimal_spec, mock_mode=True) + orchestrator.runtime_state.phase_completed = {"0": True} + + controller = DriftDetectionController( + orchestrator=orchestrator, + poll_interval_seconds=1, + auto_heal=False, + ) + + # Iteration 1: phase is healthy + healthy = True + for phase_num, handler_class in orchestrator.PHASE_HANDLERS: + if phase_num == 0: + handler_name = handler_class.__name__ + controller._update_drift_state(phase_num, handler_name, healthy) + + assert controller.drift_states[0].is_drifted is False + + # Iteration 2: phase becomes unhealthy + healthy = False + for phase_num, handler_class in orchestrator.PHASE_HANDLERS: + if phase_num == 0: + handler_name = handler_class.__name__ + controller._update_drift_state(phase_num, handler_name, healthy) + + assert controller.drift_states[0].is_drifted is True + assert controller.drift_states[0].consecutive_drift_count == 1 + + # Iteration 3: phase still unhealthy + for phase_num, handler_class in orchestrator.PHASE_HANDLERS: + if phase_num == 0: + handler_name = handler_class.__name__ + controller._update_drift_state(phase_num, handler_name, healthy) + + assert controller.drift_states[0].consecutive_drift_count == 2 + + # Iteration 4: phase recovers + healthy = True + for phase_num, handler_class in orchestrator.PHASE_HANDLERS: + if phase_num == 0: + handler_name = handler_class.__name__ + controller._update_drift_state(phase_num, handler_name, healthy) + + assert controller.drift_states[0].is_drifted is False + + @pytest.mark.asyncio + async def test_auto_healing_triggers_on_drift( + self, + minimal_spec: NetEngineSpec, + tmp_path, + monkeypatch, + ) -> None: + """Test that auto-healing is triggered when drift is detected.""" + state_file = str(tmp_path / "netengine_state.json") + monkeypatch.setenv("NETENGINE_STATE_FILE", state_file) + + orchestrator = Orchestrator(minimal_spec, mock_mode=True) + orchestrator.runtime_state.phase_completed = {"0": True} + + # Mock the orchestrator._check_prerequisites to not raise errors + orchestrator._check_prerequisites = MagicMock() + + controller = DriftDetectionController( + orchestrator=orchestrator, + poll_interval_seconds=1, + auto_heal=True, + ) + + # Create a drifted phase list + drifted_phases = [0] + + # Mock handler for self-healing + mock_handler = AsyncMock() + mock_handler.execute = AsyncMock() + mock_handler.healthcheck = AsyncMock(return_value=True) + + # Patch orchestrator.PHASE_HANDLERS to return our mock handler + with patch.object( + orchestrator, + 'PHASE_HANDLERS', + [(0, MagicMock(return_value=mock_handler))] + ): + # This would trigger self-healing in the real flow + await controller._trigger_self_healing(drifted_phases) + + # Verify execute was called during healing + mock_handler.execute.assert_called() + + @pytest.mark.asyncio + async def test_drift_history_persisted( + self, + minimal_spec: NetEngineSpec, + tmp_path, + monkeypatch, + ) -> None: + """Test that drift history is persisted to RuntimeState.""" + state_file = str(tmp_path / "netengine_state.json") + monkeypatch.setenv("NETENGINE_STATE_FILE", state_file) + + orchestrator = Orchestrator(minimal_spec, mock_mode=True) + + # Add a drift event to history + orchestrator.runtime_state.drift_history.append({ + "phase_num": 0, + "detected_at": "2025-06-26T12:00:00", + "healed_at": "2025-06-26T12:00:05", + "healing_failed": False, + "error": None, + }) + + # Save state + orchestrator.runtime_state.save() + + # Load state in a new instance + loaded_state = RuntimeState.load() + + # Verify history was persisted + assert len(loaded_state.drift_history) == 1 + assert loaded_state.drift_history[0]["phase_num"] == 0 + + @pytest.mark.asyncio + async def test_drift_controller_cancellation( + self, + minimal_spec: NetEngineSpec, + tmp_path, + monkeypatch, + ) -> None: + """Test that drift controller gracefully handles cancellation.""" + state_file = str(tmp_path / "netengine_state.json") + monkeypatch.setenv("NETENGINE_STATE_FILE", state_file) + + orchestrator = Orchestrator(minimal_spec, mock_mode=True) + + controller = DriftDetectionController( + orchestrator=orchestrator, + poll_interval_seconds=1, + auto_heal=False, + ) + + # Start the controller in a task and immediately cancel it + task = asyncio.create_task(controller.run()) + await asyncio.sleep(0.1) + task.cancel() + + try: + await task + except asyncio.CancelledError: + pass # Expected + + @pytest.mark.asyncio + async def test_event_emission_on_drift_detection( + self, + minimal_spec: NetEngineSpec, + tmp_path, + monkeypatch, + ) -> None: + """Test that drift events are emitted via pgmq.""" + state_file = str(tmp_path / "netengine_state.json") + monkeypatch.setenv("NETENGINE_STATE_FILE", state_file) + + orchestrator = Orchestrator(minimal_spec, mock_mode=True) + orchestrator.context.pgmq_client = AsyncMock() + orchestrator.context.pgmq_client.send = AsyncMock() + + controller = DriftDetectionController(orchestrator=orchestrator) + + # Emit a drift event + await controller._emit_drift_event( + phase_num=0, + handler_name="TestHandler", + event_type="drift.detected", + payload={"phase": 0}, + ) + + # Verify event was sent + orchestrator.context.pgmq_client.send.assert_called_once() + call_args = orchestrator.context.pgmq_client.send.call_args + # call_args[0][0] is queue_name, call_args[0][1] is event + event = call_args[0][1] if len(call_args[0]) > 1 else None + assert event is not None + assert event.event_type == "drift.detected" diff --git a/tests/test_drift_controller.py b/tests/test_drift_controller.py new file mode 100644 index 0000000..360d3d1 --- /dev/null +++ b/tests/test_drift_controller.py @@ -0,0 +1,298 @@ +"""Unit tests for drift detection and self-healing.""" + +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from netengine.core.drift_controller import DriftDetectionController, DriftState +from netengine.core.orchestrator import Orchestrator +from netengine.core.state import RuntimeState +from netengine.handlers.context import PhaseContext +from netengine.spec.models import NetEngineSpec + + +class TestDriftDetectionController: + """Tests for DriftDetectionController.""" + + @pytest.fixture + def mock_orchestrator(self, minimal_spec: NetEngineSpec) -> Orchestrator: + """Create a mock orchestrator.""" + orch = MagicMock(spec=Orchestrator) + orch.spec = minimal_spec + orch.runtime_state = RuntimeState() + orch.runtime_state.phase_completed = {"0": True, "1": True, "3": True} + orch.context = MagicMock(spec=PhaseContext) + orch.context.pgmq_client = MagicMock() + orch.context.pgmq_client.send = AsyncMock() + orch.PHASE_HANDLERS = [ + (0, MagicMock), + (1, MagicMock), + (3, MagicMock), + ] + return orch + + @pytest.mark.asyncio + async def test_drift_detection_initialization( + self, mock_orchestrator: Orchestrator + ) -> None: + """Test controller initialization.""" + controller = DriftDetectionController( + orchestrator=mock_orchestrator, + poll_interval_seconds=10, + max_drift_retries=2, + auto_heal=True, + ) + + assert controller.orchestrator == mock_orchestrator + assert controller.poll_interval_seconds == 10 + assert controller.max_drift_retries == 2 + assert controller.auto_heal is True + assert controller.drift_states == {} + + @pytest.mark.asyncio + async def test_check_phase_health_success( + self, mock_orchestrator: Orchestrator + ) -> None: + """Test healthcheck for a healthy phase.""" + controller = DriftDetectionController(orchestrator=mock_orchestrator) + + mock_handler = AsyncMock() + mock_handler.__class__.__name__ = "TestHandler" + mock_handler.healthcheck = AsyncMock(return_value=True) + + result = await controller._check_phase_health(0, mock_handler) + + assert result is True + mock_handler.healthcheck.assert_called_once_with(mock_orchestrator.context) + + @pytest.mark.asyncio + async def test_check_phase_health_failure( + self, mock_orchestrator: Orchestrator + ) -> None: + """Test healthcheck for a drifted phase.""" + controller = DriftDetectionController(orchestrator=mock_orchestrator) + + mock_handler = AsyncMock() + mock_handler.__class__.__name__ = "TestHandler" + mock_handler.healthcheck = AsyncMock(return_value=False) + + result = await controller._check_phase_health(0, mock_handler) + + assert result is False + assert 0 in controller.drift_states + assert controller.drift_states[0].is_drifted is True + + @pytest.mark.asyncio + async def test_check_phase_health_exception( + self, mock_orchestrator: Orchestrator + ) -> None: + """Test healthcheck that raises an exception.""" + controller = DriftDetectionController(orchestrator=mock_orchestrator) + + mock_handler = AsyncMock() + mock_handler.__class__.__name__ = "TestHandler" + mock_handler.healthcheck = AsyncMock( + side_effect=RuntimeError("health check error") + ) + + result = await controller._check_phase_health(0, mock_handler) + + assert result is False + assert 0 in controller.drift_states + assert controller.drift_states[0].is_drifted is True + + @pytest.mark.asyncio + async def test_drift_state_tracking(self, mock_orchestrator: Orchestrator) -> None: + """Test drift state tracking over multiple checks.""" + controller = DriftDetectionController(orchestrator=mock_orchestrator) + + # First check: phase is healthy + controller._update_drift_state(0, "TestHandler", True) + assert controller.drift_states[0].is_drifted is False + assert controller.drift_states[0].consecutive_drift_count == 0 + + # Second check: phase becomes unhealthy + controller._update_drift_state(0, "TestHandler", False) + assert controller.drift_states[0].is_drifted is True + assert controller.drift_states[0].consecutive_drift_count == 1 + assert controller.drift_states[0].drift_detected_at is not None + + # Third check: phase still unhealthy + controller._update_drift_state(0, "TestHandler", False) + assert controller.drift_states[0].is_drifted is True + assert controller.drift_states[0].consecutive_drift_count == 2 + + # Fourth check: phase recovers + controller._update_drift_state(0, "TestHandler", True) + assert controller.drift_states[0].is_drifted is False + assert controller.drift_states[0].consecutive_drift_count == 0 + + @pytest.mark.asyncio + async def test_drift_event_emission(self, mock_orchestrator: Orchestrator) -> None: + """Test that drift events are emitted.""" + controller = DriftDetectionController(orchestrator=mock_orchestrator) + + await controller._emit_drift_event( + phase_num=0, + handler_name="TestHandler", + event_type="drift.detected", + payload={"phase": 0, "handler": "TestHandler"}, + ) + + assert mock_orchestrator.context.pgmq_client is not None + mock_orchestrator.context.pgmq_client.send.assert_called_once() # type: ignore + + @pytest.mark.asyncio + async def test_drift_event_emission_no_pgmq( + self, mock_orchestrator: Orchestrator + ) -> None: + """Test that drift events are skipped when pgmq is unavailable.""" + mock_orchestrator.context.pgmq_client = None + controller = DriftDetectionController(orchestrator=mock_orchestrator) + + await controller._emit_drift_event( + phase_num=0, + handler_name="TestHandler", + event_type="drift.detected", + payload={"phase": 0}, + ) + + # Should not raise an error + + @pytest.mark.asyncio + async def test_heal_phase_success(self, mock_orchestrator: Orchestrator) -> None: + """Test successful phase healing.""" + controller = DriftDetectionController(orchestrator=mock_orchestrator) + + mock_handler = AsyncMock() + mock_handler.__class__.__name__ = "TestHandler" + mock_handler.execute = AsyncMock() + mock_handler.healthcheck = AsyncMock(return_value=True) + + success, changed = await controller._heal_phase(0, mock_handler) + + assert success is True + assert changed is True + mock_handler.execute.assert_called_once() + mock_handler.healthcheck.assert_called_once() + + @pytest.mark.asyncio + async def test_heal_phase_failure(self, mock_orchestrator: Orchestrator) -> None: + """Test failed phase healing.""" + controller = DriftDetectionController(orchestrator=mock_orchestrator) + + mock_handler = AsyncMock() + mock_handler.__class__.__name__ = "TestHandler" + mock_handler.execute = AsyncMock(side_effect=RuntimeError("execution failed")) + + success, changed = await controller._heal_phase(0, mock_handler) + + assert success is False + assert changed is False + + @pytest.mark.asyncio + async def test_heal_phase_healthcheck_still_fails( + self, mock_orchestrator: Orchestrator + ) -> None: + """Test when healthcheck still fails after execute().""" + controller = DriftDetectionController(orchestrator=mock_orchestrator) + + mock_handler = AsyncMock() + mock_handler.__class__.__name__ = "TestHandler" + mock_handler.execute = AsyncMock() + mock_handler.healthcheck = AsyncMock(return_value=False) + + success, changed = await controller._heal_phase(0, mock_handler) + + assert success is False + assert changed is False + + @pytest.mark.asyncio + async def test_runtime_state_persistence( + self, mock_orchestrator: Orchestrator + ) -> None: + """Test that drift state is persisted to RuntimeState.""" + # Record a drift event + mock_orchestrator.runtime_state.current_drift_phases = [0, 1] + mock_orchestrator.runtime_state.last_drift_check_at = datetime.utcnow() + + # Verify drift history can be updated + event = { + "phase_num": 0, + "detected_at": datetime.utcnow().isoformat(), + "healed_at": None, + "healing_failed": False, + "error": None, + } + mock_orchestrator.runtime_state.drift_history.append(event) + + assert len(mock_orchestrator.runtime_state.drift_history) == 1 + assert mock_orchestrator.runtime_state.drift_history[0]["phase_num"] == 0 + + @pytest.mark.asyncio + async def test_iteration_with_multiple_phases( + self, mock_orchestrator: Orchestrator + ) -> None: + """Test a drift detection iteration with multiple phases.""" + # Create mock handlers that return different health states + mock_handler_0 = AsyncMock() + mock_handler_0.__class__.__name__ = "SubstrateHandler" + mock_handler_0.healthcheck = AsyncMock(return_value=True) + + mock_handler_1 = AsyncMock() + mock_handler_1.__class__.__name__ = "DNSHandler" + mock_handler_1.healthcheck = AsyncMock(return_value=False) + + mock_handler_3 = AsyncMock() + mock_handler_3.__class__.__name__ = "PKIHandler" + mock_handler_3.healthcheck = AsyncMock(return_value=True) + + def handler_factory(handler_class): + if handler_class == mock_orchestrator.PHASE_HANDLERS[0][1]: + return mock_handler_0 + elif handler_class == mock_orchestrator.PHASE_HANDLERS[1][1]: + return mock_handler_1 + elif handler_class == mock_orchestrator.PHASE_HANDLERS[2][1]: + return mock_handler_3 + return AsyncMock() + + controller = DriftDetectionController( + orchestrator=mock_orchestrator, auto_heal=False + ) + + with patch.object( + mock_orchestrator.PHASE_HANDLERS[0][1], + '__call__', + side_effect=lambda: mock_handler_0 + ): + pass + + # Run one iteration + await controller._run_one_iteration() + + # Phase 1 (DNS) should be detected as drifted + assert 1 in controller.drift_states + assert controller.drift_states[1].is_drifted is True + + +class TestDriftState: + """Tests for DriftState dataclass.""" + + def test_drift_state_creation(self) -> None: + """Test DriftState creation.""" + now = datetime.utcnow() + state = DriftState( + phase_num=0, + handler_name="TestHandler", + last_healthcheck_at=now, + is_drifted=True, + drift_detected_at=now, + consecutive_drift_count=1, + ) + + assert state.phase_num == 0 + assert state.handler_name == "TestHandler" + assert state.is_drifted is True + assert state.consecutive_drift_count == 1 + assert state.self_healing_attempted is False