Publish curated ITWorx-ModelForge source from 990401737f96c8d46a9f7b0d84994516cadcb6a0

backend/tests/test_node_lifecycle_postgres.py

@@ -0,0 +1,328 @@
+from __future__ import annotations
+
+import importlib
+import threading
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import TimeoutError as FutureTimeoutError
+from datetime import UTC, datetime
+
+import pytest
+from sqlalchemy import event, select
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import Session
+
+pytest_plugins = ("tests.test_managed_identities_postgres",)
+
+from modelforge_api.domain.agent_protocol import AgentMetadata, HeartbeatRequest
+from modelforge_api.domain.node_lifecycle import NodeLifecycleChange, NodeLifecyclePreviewRequest
+from modelforge_api.persistence.models import ComputeNode, NodeCredential
+from modelforge_api.services.node_agent import (
+    AgentAuthorizationError,
+    NodeAgentService,
+    NodeAuthenticationEvidence,
+    secret_hash,
+)
+from modelforge_api.services.node_decommission import (
+    NodeDecommissionError,
+    NodeDecommissionService,
+)
+from modelforge_api.settings import Settings
+from tests.test_managed_identities_postgres import IdentityDatabase
+
+
+def test_agent_request_waits_for_disable_and_rechecks_locked_node(
+    identity_database: IdentityDatabase,
+) -> None:
+    """PostgreSQL proof that a pre-body auth result cannot outlive a lifecycle disable."""
+    importlib.import_module("alembic.command").upgrade(identity_database.migration(), "head")
+    engine = identity_database.runtime
+    node_id = uuid.uuid4()
+    credential_id = uuid.uuid4()
+    raw = f"mfnode_{credential_id}_lifecycle-race-secret"
+    with Session(engine) as session:
+        node = ComputeNode(
+            id=node_id,
+            key=f"postgres-race-{node_id}",
+            hostname="postgres-race-node",
+            display_name="PostgreSQL lifecycle race",
+            enabled=True,
+            status="active",
+            liveness_state="offline",
+            inventory={"source": "test"},
+            observation_source="remote_agent",
+        )
+        session.add(node)
+        session.add(
+            NodeCredential(
+                id=credential_id,
+                compute_node_id=node_id,
+                secret_hash=secret_hash(raw),
+                scope="node.publish",
+            )
+        )
+        session.commit()
+
+        preview = NodeDecommissionService(session).preview_enabled_change(
+            node_id, NodeLifecyclePreviewRequest(target_enabled=False)
+        )
+        change = NodeLifecycleChange(
+            target_enabled=False,
+            expected_generation=preview.node_generation,
+            dependency_digest=preview.dependency_digest,
+            operator="postgres-race-test",
+            reason="Prove that disable serializes with agent authentication.",
+            confirmation=preview.persisted_identity,
+        )
+
+    request_session = Session(engine, expire_on_commit=False)
+    agent_service = NodeAgentService(request_session, Settings(_env_file=None))
+    (_credential, stale_node), evidence = agent_service.authenticate_with_evidence(
+        f"Bearer {raw}"
+    )
+    request_session.commit()
+    assert stale_node.enabled is True
+
+    lifecycle_has_locks = threading.Event()
+    allow_lifecycle_to_commit = threading.Event()
+    agent_lock_attempted = threading.Event()
+
+    with Session(engine) as lifecycle_session:
+        lifecycle_service = NodeDecommissionService(lifecycle_session)
+        lock_dependencies = lifecycle_service._lock_dependency_tables
+
+        def hold_lifecycle_locks() -> None:
+            lock_dependencies()
+            lifecycle_has_locks.set()
+            if not allow_lifecycle_to_commit.wait(timeout=20):
+                raise RuntimeError("timed out holding lifecycle locks in race test")
+
+        lifecycle_service._lock_dependency_tables = hold_lifecycle_locks  # type: ignore[method-assign]
+
+        @event.listens_for(engine, "before_cursor_execute")
+        def observe_agent_lock_attempt(
+            _connection: object,
+            _cursor: object,
+            statement: str,
+            _parameters: object,
+            _context: object,
+            _executemany: bool,
+        ) -> None:
+            normalized = statement.lower()
+            if (
+                threading.current_thread().name.startswith("agent-request")
+                and "compute_nodes" in normalized
+                and "for update" in normalized
+            ):
+                agent_lock_attempted.set()
+
+        try:
+            with ThreadPoolExecutor(max_workers=1, thread_name_prefix="lifecycle") as lifecycle_pool:
+                lifecycle_future = lifecycle_pool.submit(
+                    lifecycle_service.apply_enabled_change, node_id, change
+                )
+                assert lifecycle_has_locks.wait(timeout=10), "lifecycle never acquired its locks"
+                with ThreadPoolExecutor(max_workers=1, thread_name_prefix="agent-request") as agent_pool:
+                    agent_future = agent_pool.submit(
+                        _publish_authenticated_heartbeat,
+                        agent_service,
+                        evidence,
+                        raw,
+                        stale_node.key,
+                    )
+                    assert agent_lock_attempted.wait(timeout=10), "agent never requested its node lock"
+                    with pytest.raises(FutureTimeoutError):
+                        agent_future.result(timeout=0.25)
+
+                    allow_lifecycle_to_commit.set()
+                    lifecycle_future.result(timeout=30)
+                    assert agent_future.result(timeout=30) == "rejected"
+        finally:
+            allow_lifecycle_to_commit.set()
+            event.remove(engine, "before_cursor_execute", observe_agent_lock_attempt)
+            request_session.close()
+
+    with Session(engine) as session:
+        persisted = session.scalar(select(ComputeNode).where(ComputeNode.id == node_id))
+        assert persisted is not None
+        assert persisted.enabled is False
+        assert persisted.liveness_state == "disabled"
+        assert persisted.generation == change.expected_generation + 1
+
+
+def test_heartbeat_wins_node_lock_and_invalidates_lifecycle_preview(
+    identity_database: IdentityDatabase,
+) -> None:
+    """A heartbeat serialized before disable advances the CAS generation."""
+    importlib.import_module("alembic.command").upgrade(identity_database.migration(), "head")
+    engine = identity_database.runtime
+    node_id = uuid.uuid4()
+    credential_id = uuid.uuid4()
+    raw = f"mfnode_{credential_id}_heartbeat-first-race-secret"
+    with Session(engine) as session:
+        node = ComputeNode(
+            id=node_id,
+            key=f"postgres-heartbeat-race-{node_id}",
+            hostname="postgres-heartbeat-race-node",
+            display_name="PostgreSQL heartbeat race",
+            enabled=True,
+            status="active",
+            liveness_state="offline",
+            inventory={"source": "test"},
+            observation_source="remote_agent",
+        )
+        session.add(node)
+        session.add(
+            NodeCredential(
+                id=credential_id,
+                compute_node_id=node_id,
+                secret_hash=secret_hash(raw),
+                scope="node.publish",
+            )
+        )
+        session.commit()
+        lifecycle = NodeDecommissionService(session)
+        preview = lifecycle.preview_enabled_change(
+            node_id, NodeLifecyclePreviewRequest(target_enabled=False)
+        )
+        change = NodeLifecycleChange(
+            target_enabled=False,
+            expected_generation=preview.node_generation,
+            dependency_digest=preview.dependency_digest,
+            operator="postgres-race-test",
+            reason="Prove that fresh heartbeat evidence invalidates this preview.",
+            confirmation=preview.persisted_identity,
+        )
+
+    agent_session = Session(engine, expire_on_commit=False)
+    agent_service = NodeAgentService(agent_session, Settings(_env_file=None))
+    (_credential, stale_node), evidence = agent_service.authenticate_with_evidence(
+        f"Bearer {raw}"
+    )
+    agent_session.commit()
+    heartbeat_has_node_lock = threading.Event()
+    allow_heartbeat_to_commit = threading.Event()
+    lifecycle_lock_attempted = threading.Event()
+    original_commit = agent_session.commit
+
+    def hold_heartbeat_commit() -> None:
+        heartbeat_has_node_lock.set()
+        if not allow_heartbeat_to_commit.wait(timeout=20):
+            raise RuntimeError("timed out holding the heartbeat node lock")
+        original_commit()
+
+    agent_session.commit = hold_heartbeat_commit  # type: ignore[method-assign]
+
+    @event.listens_for(engine, "before_cursor_execute")
+    def observe_lifecycle_lock_attempt(
+        _connection: object,
+        _cursor: object,
+        statement: str,
+        _parameters: object,
+        _context: object,
+        _executemany: bool,
+    ) -> None:
+        normalized = statement.lower()
+        if (
+            threading.current_thread().name.startswith("lifecycle-after-heartbeat")
+            and "compute_nodes" in normalized
+            and "for update" in normalized
+        ):
+            lifecycle_lock_attempted.set()
+
+    request = HeartbeatRequest(
+        identity_key=stale_node.key,
+        metadata=AgentMetadata(
+            agent_version="1.3.0-rc.1",
+            protocol_version=1,
+            supported_capabilities=["hardware.inventory"],
+            started_at=datetime.now(UTC),
+        ),
+        observed_at=datetime.now(UTC),
+    )
+
+    try:
+        with ThreadPoolExecutor(max_workers=1, thread_name_prefix="heartbeat") as agent_pool:
+            heartbeat_future = agent_pool.submit(
+                _publish_and_hold_authenticated_heartbeat,
+                agent_service,
+                evidence,
+                raw,
+                request,
+            )
+            assert heartbeat_has_node_lock.wait(timeout=10), "heartbeat never acquired node lock"
+            with ThreadPoolExecutor(
+                max_workers=1, thread_name_prefix="lifecycle-after-heartbeat"
+            ) as lifecycle_pool:
+                lifecycle_future = lifecycle_pool.submit(
+                    _apply_stale_lifecycle_change, engine, node_id, change
+                )
+                assert lifecycle_lock_attempted.wait(timeout=10), "lifecycle never requested node lock"
+                with pytest.raises(FutureTimeoutError):
+                    lifecycle_future.result(timeout=0.25)
+                allow_heartbeat_to_commit.set()
+                heartbeat_future.result(timeout=30)
+                assert lifecycle_future.result(timeout=30) == "node_generation_conflict"
+    finally:
+        allow_heartbeat_to_commit.set()
+        event.remove(engine, "before_cursor_execute", observe_lifecycle_lock_attempt)
+        agent_session.close()
+
+    with Session(engine) as session:
+        persisted = session.scalar(select(ComputeNode).where(ComputeNode.id == node_id))
+        assert persisted is not None
+        assert persisted.enabled is True
+        assert persisted.liveness_state == "online"
+        assert persisted.last_heartbeat_at is not None
+        assert persisted.generation == change.expected_generation + 1
+
+
+def _publish_and_hold_authenticated_heartbeat(
+    service: NodeAgentService,
+    evidence: NodeAuthenticationEvidence,
+    raw_credential: str,
+    request: HeartbeatRequest,
+) -> None:
+    _credential, node = service.reuse_authentication(evidence, f"Bearer {raw_credential}")
+    service.heartbeat(node, request)
+
+
+def _apply_stale_lifecycle_change(
+    engine: Engine, node_id: uuid.UUID, change: NodeLifecycleChange
+) -> str:
+    with Session(engine) as session:
+        try:
+            NodeDecommissionService(session).apply_enabled_change(node_id, change)
+        except NodeDecommissionError as exc:
+            if exc.code != "node_generation_conflict":
+                raise
+            return exc.code
+    return "applied"
+
+
+def _publish_authenticated_heartbeat(
+    service: NodeAgentService,
+    evidence: NodeAuthenticationEvidence,
+    raw_credential: str,
+    identity_key: str,
+) -> str:
+    try:
+        _credential, node = service.reuse_authentication(
+            evidence, f"Bearer {raw_credential}"
+        )
+        service.heartbeat(
+            node,
+            HeartbeatRequest(
+                identity_key=identity_key,
+                metadata=AgentMetadata(
+                    agent_version="1.3.0-rc.1",
+                    protocol_version=1,
+                    supported_capabilities=["hardware.inventory"],
+                    started_at=datetime.now(UTC),
+                ),
+                observed_at=datetime.now(UTC),
+            ),
+        )
+    except AgentAuthorizationError:
+        return "rejected"
+    return "published"