142 lines
6.1 KiB
Python
142 lines
6.1 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
from django.core.exceptions import ValidationError
|
|
from django.db import IntegrityError
|
|
|
|
from control_plane.agents.models import ProgenySignal
|
|
from graph.models import ExecutionGraphDefinition, ExecutionGraphVersion, ExecutionGraphVersionStatus, GraphRun, GraphRunStatus
|
|
from graph.registry import NodeHandlerRegistry, NodeResult
|
|
from graph.spec import ExecutionGraphSpec, outgoing_edges
|
|
from graph.task_execution import task_execution_graph_v1
|
|
|
|
|
|
class FixtureNode:
|
|
node_type = "fixture"
|
|
idempotent = True
|
|
replay_safe = True
|
|
destructive = False
|
|
|
|
def run(self, context: object) -> NodeResult:
|
|
return NodeResult("COMPLETE", "success", {"ok": True})
|
|
|
|
|
|
def test_task_execution_graph_v1_is_serializable_and_valid() -> None:
|
|
spec = task_execution_graph_v1()
|
|
|
|
payload = spec.to_dict()
|
|
restored = ExecutionGraphSpec.from_dict(payload)
|
|
restored.validate()
|
|
|
|
assert restored.name == "task_execution"
|
|
assert restored.version == 1
|
|
assert restored.entry == "claim_task"
|
|
assert restored.nodes["coder"].metadata["contains_internal_tool_loop"] is True
|
|
assert {edge.condition for edge in outgoing_edges(restored, "review")} == {"PASS", "REWORK_REQUIRED", "REJECTED"}
|
|
assert {edge.condition for edge in outgoing_edges(restored, "retry_or_fail")} == {"retry_available", "retry_exhausted"}
|
|
|
|
|
|
def test_node_registry_resolves_handlers_without_arbitrary_execution() -> None:
|
|
registry = NodeHandlerRegistry()
|
|
registry.register(FixtureNode())
|
|
|
|
assert registry.get("fixture").run(object()).output_metadata == {"ok": True}
|
|
with pytest.raises(KeyError):
|
|
registry.get("missing")
|
|
|
|
|
|
def test_graph_models_persist_versioned_run_state() -> None:
|
|
spec = task_execution_graph_v1()
|
|
definition = ExecutionGraphDefinition.objects.create(
|
|
name=spec.name,
|
|
graph_type=spec.graph_type,
|
|
description="Task execution graph",
|
|
)
|
|
version = ExecutionGraphVersion.objects.create(
|
|
graph=definition,
|
|
version=spec.version,
|
|
status=ExecutionGraphVersionStatus.CHAMPION,
|
|
graph_spec=spec.to_dict(),
|
|
metadata={"immutable_after_use": True},
|
|
)
|
|
|
|
run = GraphRun.objects.create(
|
|
execution_graph_version=version,
|
|
status=GraphRunStatus.PENDING,
|
|
current_node=spec.entry,
|
|
metadata={"node_count": len(spec.nodes)},
|
|
)
|
|
|
|
assert run.execution_graph_version.graph.name == "task_execution"
|
|
assert run.execution_graph_version.status == ExecutionGraphVersionStatus.CHAMPION
|
|
assert run.metadata["node_count"] == 12
|
|
|
|
|
|
def test_champion_graph_version_spec_is_immutable() -> None:
|
|
spec = task_execution_graph_v1()
|
|
definition = ExecutionGraphDefinition.objects.create(name="immutable_task_execution", graph_type=spec.graph_type)
|
|
version = ExecutionGraphVersion.objects.create(graph=definition, version=1, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())
|
|
|
|
version.graph_spec = {**version.graph_spec, "entry": "different"}
|
|
|
|
with pytest.raises(ValidationError):
|
|
version.save()
|
|
|
|
|
|
def test_graph_changes_require_a_new_version() -> None:
|
|
spec = task_execution_graph_v1()
|
|
definition = ExecutionGraphDefinition.objects.create(name="versioned_task_execution", graph_type=spec.graph_type)
|
|
v1 = ExecutionGraphVersion.objects.create(graph=definition, version=1, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())
|
|
changed_spec = spec.to_dict()
|
|
changed_spec["metadata"] = {**changed_spec["metadata"], "change": "new node planned"}
|
|
|
|
v2 = ExecutionGraphVersion.objects.create(graph=definition, version=2, status=ExecutionGraphVersionStatus.CHALLENGER, graph_spec=changed_spec)
|
|
|
|
assert v1.graph_spec != v2.graph_spec
|
|
assert v2.status == ExecutionGraphVersionStatus.CHALLENGER
|
|
|
|
|
|
def test_only_one_champion_version_per_graph_definition() -> None:
|
|
spec = task_execution_graph_v1()
|
|
definition = ExecutionGraphDefinition.objects.create(name="single_champion_task_execution", graph_type=spec.graph_type)
|
|
ExecutionGraphVersion.objects.create(graph=definition, version=1, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())
|
|
|
|
with pytest.raises((IntegrityError, ValidationError)):
|
|
ExecutionGraphVersion.objects.create(graph=definition, version=2, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())
|
|
|
|
|
|
def test_commit_lineage_preserves_exact_graph_version() -> None:
|
|
spec = task_execution_graph_v1()
|
|
definition = ExecutionGraphDefinition.objects.create(name="lineage_preserved_task_execution", graph_type=spec.graph_type)
|
|
v1 = ExecutionGraphVersion.objects.create(graph=definition, version=1, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())
|
|
v2 = ExecutionGraphVersion.objects.create(graph=definition, version=2, status=ExecutionGraphVersionStatus.CHALLENGER, graph_spec=spec.to_dict())
|
|
run = GraphRun.objects.create(execution_graph_version=v1, status=GraphRunStatus.COMPLETE, current_node="complete")
|
|
|
|
v2.status = ExecutionGraphVersionStatus.RETIRED
|
|
v2.save()
|
|
run.refresh_from_db()
|
|
|
|
assert run.execution_graph_version == v1
|
|
assert run.execution_graph_version.version == 1
|
|
|
|
|
|
def test_progeny_signal_can_reference_graph_lineage() -> None:
|
|
spec = task_execution_graph_v1()
|
|
definition = ExecutionGraphDefinition.objects.create(name="lineage_graph", graph_type=spec.graph_type)
|
|
version = ExecutionGraphVersion.objects.create(graph=definition, version=1, status=ExecutionGraphVersionStatus.CHAMPION, graph_spec=spec.to_dict())
|
|
run = GraphRun.objects.create(execution_graph_version=version, current_node=spec.entry)
|
|
node_run = run.node_runs.create(node_id="coder", node_type="coder", visit_index=1)
|
|
|
|
signal = ProgenySignal.objects.create(
|
|
source="runtime",
|
|
severity="high",
|
|
failure_category="MODEL_OUTPUT_INVALID",
|
|
summary="Malformed model JSON",
|
|
graph_run=run,
|
|
graph_node_run=node_run,
|
|
execution_graph_version=version,
|
|
)
|
|
|
|
assert signal.graph_run == run
|
|
assert signal.graph_node_run == node_run
|
|
assert signal.execution_graph_version == version
|