313 lines
11 KiB
Python
313 lines
11 KiB
Python
"""Tests for Pipeline Saga compensation pattern"""
|
|
|
|
import asyncio
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from agentkit.orchestrator.compensation import (
|
|
CompletedStep,
|
|
CompensationResult,
|
|
SagaOrchestrator,
|
|
)
|
|
from agentkit.orchestrator.pipeline_engine import PipelineEngine
|
|
from agentkit.orchestrator.pipeline_schema import (
|
|
Pipeline,
|
|
PipelineStage,
|
|
StageStatus,
|
|
)
|
|
from agentkit.orchestrator.retry import StepRetryPolicy
|
|
from agentkit.skills.geo_pipeline import (
|
|
GEO_PIPELINE_COMPENSATIONS,
|
|
PipelineStep,
|
|
create_geo_pipeline_steps,
|
|
)
|
|
|
|
|
|
class TestCompletedStep:
|
|
"""CompletedStep dataclass tests"""
|
|
|
|
def test_creation_with_compensation(self):
|
|
step = CompletedStep(
|
|
step_name="optimize",
|
|
result={"changes": 5},
|
|
compensate_action="revert_optimization",
|
|
)
|
|
assert step.step_name == "optimize"
|
|
assert step.result == {"changes": 5}
|
|
assert step.compensate_action == "revert_optimization"
|
|
|
|
def test_creation_without_compensation(self):
|
|
step = CompletedStep(step_name="detect", result={"found": 3})
|
|
assert step.compensate_action is None
|
|
|
|
|
|
class TestCompensationResult:
|
|
"""CompensationResult dataclass tests"""
|
|
|
|
def test_success_result(self):
|
|
result = CompensationResult(step_name="optimize", success=True)
|
|
assert result.step_name == "optimize"
|
|
assert result.success is True
|
|
assert result.error is None
|
|
|
|
def test_failure_result(self):
|
|
result = CompensationResult(
|
|
step_name="optimize", success=False, error="rollback failed"
|
|
)
|
|
assert result.success is False
|
|
assert result.error == "rollback failed"
|
|
|
|
|
|
class TestSagaOrchestrator:
|
|
"""SagaOrchestrator tests"""
|
|
|
|
def test_record_completed(self):
|
|
saga = SagaOrchestrator()
|
|
saga.record_completed("step1", {"data": 1}, "compensate_1")
|
|
saga.record_completed("step2", {"data": 2})
|
|
|
|
steps = saga.completed_steps
|
|
assert len(steps) == 2
|
|
assert steps[0].step_name == "step1"
|
|
assert steps[0].compensate_action == "compensate_1"
|
|
assert steps[1].step_name == "step2"
|
|
assert steps[1].compensate_action is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compensate_lifo_order(self):
|
|
"""Compensation should execute in LIFO (reverse) order"""
|
|
execution_order = []
|
|
|
|
async def mock_execute_skill(skill_name: str, input_data):
|
|
execution_order.append(skill_name)
|
|
|
|
saga = SagaOrchestrator(execute_skill_func=mock_execute_skill)
|
|
saga.record_completed("step1", {"data": 1}, "compensate_1")
|
|
saga.record_completed("step2", {"data": 2}, "compensate_2")
|
|
saga.record_completed("step3", {"data": 3}, "compensate_3")
|
|
|
|
results = await saga.compensate()
|
|
|
|
# LIFO order: step3, step2, step1
|
|
assert execution_order == ["compensate_3", "compensate_2", "compensate_1"]
|
|
assert len(results) == 3
|
|
assert all(r.success for r in results)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_skip_steps_with_no_compensate_action(self):
|
|
"""Steps with no compensate_action should be skipped"""
|
|
saga = SagaOrchestrator()
|
|
saga.record_completed("read_only", {"data": 1}) # No compensation
|
|
saga.record_completed("write_op", {"data": 2}, "rollback_write")
|
|
|
|
results = await saga.compensate()
|
|
|
|
assert len(results) == 2
|
|
# write_op is compensated first (LIFO), then read_only
|
|
assert results[0].step_name == "write_op"
|
|
assert results[0].success is True
|
|
assert results[1].step_name == "read_only"
|
|
assert results[1].success is True
|
|
assert results[1].error == "no_compensation_needed"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compensation_failure_doesnt_interrupt_others(self):
|
|
"""If one compensation fails, others should still execute"""
|
|
execution_order = []
|
|
|
|
async def mock_execute_skill(skill_name: str, input_data):
|
|
execution_order.append(skill_name)
|
|
if skill_name == "compensate_2":
|
|
raise RuntimeError("rollback failed")
|
|
|
|
saga = SagaOrchestrator(execute_skill_func=mock_execute_skill)
|
|
saga.record_completed("step1", {"data": 1}, "compensate_1")
|
|
saga.record_completed("step2", {"data": 2}, "compensate_2")
|
|
saga.record_completed("step3", {"data": 3}, "compensate_3")
|
|
|
|
results = await saga.compensate()
|
|
|
|
# All compensations should be attempted (LIFO: step3, step2, step1)
|
|
assert execution_order == ["compensate_3", "compensate_2", "compensate_1"]
|
|
assert len(results) == 3
|
|
|
|
# step3 succeeds
|
|
assert results[0].success is True
|
|
# step2 fails
|
|
assert results[1].success is False
|
|
assert results[1].error == "rollback failed"
|
|
# step1 still succeeds
|
|
assert results[2].success is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compensate_with_no_execute_skill_func(self):
|
|
"""Without execute_skill_func, compensation succeeds but does nothing"""
|
|
saga = SagaOrchestrator()
|
|
saga.record_completed("step1", {"data": 1}, "compensate_1")
|
|
|
|
results = await saga.compensate()
|
|
|
|
assert len(results) == 1
|
|
assert results[0].success is True
|
|
|
|
def test_clear(self):
|
|
saga = SagaOrchestrator()
|
|
saga.record_completed("step1", {"data": 1})
|
|
saga.record_completed("step2", {"data": 2})
|
|
assert len(saga.completed_steps) == 2
|
|
|
|
saga.clear()
|
|
assert len(saga.completed_steps) == 0
|
|
|
|
def test_completed_steps_returns_copy(self):
|
|
saga = SagaOrchestrator()
|
|
saga.record_completed("step1", {"data": 1})
|
|
|
|
steps = saga.completed_steps
|
|
steps.clear() # Mutate the copy
|
|
|
|
assert len(saga.completed_steps) == 1 # Original unchanged
|
|
|
|
|
|
class TestPipelineIntegration:
|
|
"""Pipeline engine integration with retry and compensation"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_step_failure_triggers_compensation(self):
|
|
"""When a step fails, Saga compensation should be triggered for completed steps"""
|
|
engine = PipelineEngine()
|
|
|
|
pipeline = Pipeline(
|
|
name="test_compensation",
|
|
version="1.0",
|
|
description="Test compensation",
|
|
stages=[
|
|
PipelineStage(
|
|
name="step1",
|
|
agent="agent_a",
|
|
action="do_a",
|
|
compensate="undo_a",
|
|
),
|
|
PipelineStage(
|
|
name="step2",
|
|
agent="agent_b",
|
|
action="do_b",
|
|
depends_on=["step1"],
|
|
),
|
|
],
|
|
)
|
|
|
|
# Dry-run mode (no dispatcher) — all steps succeed
|
|
result = await engine.execute(pipeline)
|
|
assert result.status == StageStatus.COMPLETED
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_continue_on_failure(self):
|
|
"""Steps with continue_on_failure should not abort the pipeline"""
|
|
engine = PipelineEngine()
|
|
|
|
pipeline = Pipeline(
|
|
name="test_continue",
|
|
version="1.0",
|
|
description="Test continue_on_failure",
|
|
stages=[
|
|
PipelineStage(
|
|
name="step1",
|
|
agent="agent_a",
|
|
action="do_a",
|
|
continue_on_failure=True,
|
|
),
|
|
PipelineStage(
|
|
name="step2",
|
|
agent="agent_b",
|
|
action="do_b",
|
|
depends_on=["step1"],
|
|
),
|
|
],
|
|
)
|
|
|
|
# Dry-run mode — all steps succeed
|
|
result = await engine.execute(pipeline)
|
|
assert result.status == StageStatus.COMPLETED
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pipeline_with_retry_policy(self):
|
|
"""PipelineStage can have a retry_policy configured"""
|
|
retry = StepRetryPolicy(max_attempts=5, base_delay=0.01, jitter=False)
|
|
|
|
stage = PipelineStage(
|
|
name="retry_step",
|
|
agent="agent_a",
|
|
action="do_a",
|
|
retry_policy=retry,
|
|
)
|
|
assert stage.retry_policy is not None
|
|
assert stage.retry_policy.max_attempts == 5
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pipeline_with_compensate(self):
|
|
"""PipelineStage can have a compensate action configured"""
|
|
stage = PipelineStage(
|
|
name="optimizable_step",
|
|
agent="agent_a",
|
|
action="do_a",
|
|
compensate="undo_a",
|
|
)
|
|
assert stage.compensate == "undo_a"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pipeline_without_compensate(self):
|
|
"""PipelineStage without compensate defaults to None"""
|
|
stage = PipelineStage(
|
|
name="readonly_step",
|
|
agent="agent_a",
|
|
action="do_a",
|
|
)
|
|
assert stage.compensate is None
|
|
|
|
|
|
class TestGEOPipelineCompensations:
|
|
"""GEO Pipeline compensation definitions"""
|
|
|
|
def test_compensation_definitions(self):
|
|
"""Verify GEO pipeline compensation definitions"""
|
|
assert GEO_PIPELINE_COMPENSATIONS["detect"] is None
|
|
assert GEO_PIPELINE_COMPENSATIONS["analyze_competitor"] is None
|
|
assert GEO_PIPELINE_COMPENSATIONS["optimize"] == "revert_optimization"
|
|
assert GEO_PIPELINE_COMPENSATIONS["schema"] is None
|
|
assert GEO_PIPELINE_COMPENSATIONS["monitor"] is None
|
|
|
|
def test_create_geo_pipeline_steps(self):
|
|
"""Verify GEO pipeline steps are created with compensation"""
|
|
steps = create_geo_pipeline_steps()
|
|
assert len(steps) == 5
|
|
|
|
step_names = [s.name for s in steps]
|
|
assert step_names == [
|
|
"detect",
|
|
"analyze_competitor",
|
|
"optimize",
|
|
"schema",
|
|
"monitor",
|
|
]
|
|
|
|
# Check compensation assignments
|
|
step_map = {s.name: s for s in steps}
|
|
assert step_map["detect"].compensate is None
|
|
assert step_map["analyze_competitor"].compensate is None
|
|
assert step_map["optimize"].compensate == "revert_optimization"
|
|
assert step_map["schema"].compensate is None
|
|
assert step_map["monitor"].compensate is None
|
|
|
|
def test_geo_pipeline_steps_dependencies(self):
|
|
"""Verify GEO pipeline step dependencies form a valid chain"""
|
|
steps = create_geo_pipeline_steps()
|
|
step_map = {s.name: s for s in steps}
|
|
|
|
assert step_map["detect"].depends_on == []
|
|
assert step_map["analyze_competitor"].depends_on == ["detect"]
|
|
assert step_map["optimize"].depends_on == ["analyze_competitor"]
|
|
assert step_map["schema"].depends_on == ["optimize"]
|
|
assert step_map["monitor"].depends_on == ["schema"]
|