110 lines
3.2 KiB
Python
110 lines
3.2 KiB
Python
"""Tests for Pipeline system"""
|
|
|
|
import pytest
|
|
|
|
from agentkit.orchestrator.pipeline_schema import Pipeline, PipelineStage
|
|
from agentkit.orchestrator.pipeline_engine import PipelineEngine
|
|
from agentkit.orchestrator.pipeline_loader import PipelineLoader
|
|
|
|
|
|
def test_pipeline_schema():
|
|
stage = PipelineStage(
|
|
name="step1",
|
|
agent="agent_a",
|
|
action="process",
|
|
depends_on=[],
|
|
inputs={"data": "${input}"},
|
|
)
|
|
pipeline = Pipeline(
|
|
name="test_pipeline",
|
|
version="1.0.0",
|
|
description="Test",
|
|
stages=[stage],
|
|
variables={"input": "hello"},
|
|
)
|
|
assert len(pipeline.stages) == 1
|
|
assert pipeline.variables["input"] == "hello"
|
|
|
|
|
|
def test_topological_group():
|
|
stages = [
|
|
PipelineStage(name="a", agent="agent1", action="do_a", depends_on=[]),
|
|
PipelineStage(name="b", agent="agent2", action="do_b", depends_on=["a"]),
|
|
PipelineStage(name="c", agent="agent3", action="do_c", depends_on=["a"]),
|
|
PipelineStage(name="d", agent="agent4", action="do_d", depends_on=["b", "c"]),
|
|
]
|
|
|
|
groups = PipelineEngine._topological_group(stages)
|
|
assert len(groups) == 3 # [a], [b, c], [d]
|
|
assert groups[0][0].name == "a"
|
|
assert {s.name for s in groups[1]} == {"b", "c"}
|
|
assert groups[2][0].name == "d"
|
|
|
|
|
|
def test_topological_group_circular():
|
|
stages = [
|
|
PipelineStage(name="a", agent="agent1", action="do_a", depends_on=["b"]),
|
|
PipelineStage(name="b", agent="agent2", action="do_b", depends_on=["a"]),
|
|
]
|
|
|
|
with pytest.raises(ValueError, match="Circular dependency"):
|
|
PipelineEngine._topological_group(stages)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pipeline_dry_run():
|
|
pipeline = Pipeline(
|
|
name="test",
|
|
version="1.0.0",
|
|
description="Test",
|
|
stages=[
|
|
PipelineStage(name="step1", agent="agent1", action="process", inputs={"x": "1"}),
|
|
PipelineStage(name="step2", agent="agent2", action="analyze", depends_on=["step1"], inputs={"y": "2"}),
|
|
],
|
|
)
|
|
|
|
engine = PipelineEngine(dispatcher=None) # dry-run mode
|
|
result = await engine.execute(pipeline)
|
|
assert result.status.value == "completed"
|
|
|
|
|
|
def test_yaml_loader():
|
|
yaml_content = """
|
|
name: test_pipeline
|
|
version: "1.0"
|
|
description: Test pipeline
|
|
stages:
|
|
- name: step1
|
|
agent: agent1
|
|
action: process
|
|
inputs:
|
|
data: hello
|
|
- name: step2
|
|
agent: agent2
|
|
action: analyze
|
|
depends_on: [step1]
|
|
inputs:
|
|
result: ${{step1_output}}
|
|
variables:
|
|
input: world
|
|
"""
|
|
loader = PipelineLoader()
|
|
pipeline = loader.load_from_yaml(yaml_content, "test")
|
|
assert pipeline.name == "test_pipeline"
|
|
assert len(pipeline.stages) == 2
|
|
assert pipeline.stages[1].depends_on == ["step1"]
|
|
|
|
|
|
def test_validate_dag():
|
|
stages = [
|
|
PipelineStage(name="a", agent="a1", action="do", depends_on=[]),
|
|
PipelineStage(name="b", agent="a2", action="do", depends_on=["a"]),
|
|
]
|
|
assert PipelineLoader.validate_dag(stages) is True
|
|
|
|
circular = [
|
|
PipelineStage(name="a", agent="a1", action="do", depends_on=["b"]),
|
|
PipelineStage(name="b", agent="a2", action="do", depends_on=["a"]),
|
|
]
|
|
assert PipelineLoader.validate_dag(circular) is False
|