232 lines
8.3 KiB
Python
232 lines
8.3 KiB
Python
"""Tests for GEOPipeline"""
|
|
|
|
import pytest
|
|
|
|
from agentkit.skills.geo_pipeline import (
|
|
GEOPipeline,
|
|
PipelineStep,
|
|
PipelineStepResult,
|
|
PipelineResult,
|
|
)
|
|
|
|
|
|
class MockAgent:
|
|
"""Mock Agent for pipeline testing"""
|
|
|
|
def __init__(self, name: str, output_data: dict | None = None):
|
|
self.name = name
|
|
self.agent_type = "mock"
|
|
self._output_data = output_data or {"result": f"output from {name}"}
|
|
|
|
async def execute(self, task):
|
|
from agentkit.core.protocol import TaskResult, TaskStatus
|
|
from datetime import datetime, timezone
|
|
now = datetime.now(timezone.utc)
|
|
return TaskResult(
|
|
task_id=task.task_id,
|
|
agent_name=self.name,
|
|
status=TaskStatus.COMPLETED,
|
|
output_data=self._output_data,
|
|
error_message=None,
|
|
started_at=now,
|
|
completed_at=now,
|
|
)
|
|
|
|
|
|
class MockAgentPool:
|
|
"""Mock AgentPool"""
|
|
|
|
def __init__(self, agents: dict[str, MockAgent] | None = None):
|
|
self._agents = agents or {}
|
|
|
|
def get_agent(self, name: str):
|
|
return self._agents.get(name)
|
|
|
|
def list_agents(self):
|
|
return [{"name": a.name, "agent_type": a.agent_type} for a in self._agents.values()]
|
|
|
|
|
|
class TestGEOPipeline:
|
|
"""GEOPipeline unit tests"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sequential_pipeline(self):
|
|
"""Sequential steps should execute in order"""
|
|
steps = [
|
|
PipelineStep(name="step1", skill="skill_a"),
|
|
PipelineStep(name="step2", skill="skill_b", depends_on=["step1"]),
|
|
]
|
|
pool = MockAgentPool({
|
|
"skill_a": MockAgent("skill_a", {"data": "result_a"}),
|
|
"skill_b": MockAgent("skill_b", {"data": "result_b"}),
|
|
})
|
|
pipeline = GEOPipeline(name="test", steps=steps, agent_pool=pool)
|
|
|
|
result = await pipeline.execute({"query": "test"})
|
|
assert result.success
|
|
assert len(result.steps) == 2
|
|
assert result.steps[0].status == "success"
|
|
assert result.steps[1].status == "success"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parallel_steps(self):
|
|
"""Steps without dependencies should execute in parallel"""
|
|
steps = [
|
|
PipelineStep(name="step1", skill="skill_a"),
|
|
PipelineStep(name="step2", skill="skill_b"),
|
|
]
|
|
pool = MockAgentPool({
|
|
"skill_a": MockAgent("skill_a", {"data": "a"}),
|
|
"skill_b": MockAgent("skill_b", {"data": "b"}),
|
|
})
|
|
pipeline = GEOPipeline(name="test", steps=steps, agent_pool=pool)
|
|
|
|
result = await pipeline.execute({"query": "test"})
|
|
assert result.success
|
|
assert len(result.steps) == 2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_dag_execution(self):
|
|
"""DAG with mixed parallel/sequential steps"""
|
|
steps = [
|
|
PipelineStep(name="detect", skill="skill_a"),
|
|
PipelineStep(name="analyze_1", skill="skill_b", depends_on=["detect"]),
|
|
PipelineStep(name="analyze_2", skill="skill_c", depends_on=["detect"]),
|
|
PipelineStep(name="optimize", skill="skill_d", depends_on=["analyze_1", "analyze_2"]),
|
|
]
|
|
pool = MockAgentPool({
|
|
"skill_a": MockAgent("skill_a", {"citations": 5}),
|
|
"skill_b": MockAgent("skill_b", {"competitor": "data"}),
|
|
"skill_c": MockAgent("skill_c", {"trend": "up"}),
|
|
"skill_d": MockAgent("skill_d", {"optimized": True}),
|
|
})
|
|
pipeline = GEOPipeline(name="test", steps=steps, agent_pool=pool)
|
|
|
|
result = await pipeline.execute({"brand": "TestBrand"})
|
|
assert result.success
|
|
assert len(result.steps) == 4
|
|
|
|
# Check execution groups
|
|
groups = pipeline._build_execution_groups()
|
|
assert len(groups) == 3 # [detect], [analyze_1, analyze_2], [optimize]
|
|
assert "detect" in groups[0]
|
|
assert set(groups[1]) == {"analyze_1", "analyze_2"}
|
|
assert groups[2] == ["optimize"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_step_failure(self):
|
|
"""Failed step should be recorded"""
|
|
class FailingAgent:
|
|
name = "skill_a"
|
|
agent_type = "mock"
|
|
async def execute(self, task):
|
|
raise RuntimeError("Agent failed")
|
|
|
|
steps = [PipelineStep(name="step1", skill="skill_a")]
|
|
pool = MockAgentPool({"skill_a": FailingAgent()})
|
|
pipeline = GEOPipeline(name="test", steps=steps, agent_pool=pool)
|
|
|
|
result = await pipeline.execute({"query": "test"})
|
|
assert not result.success
|
|
assert result.steps[0].status == "failed"
|
|
assert "Agent failed" in result.steps[0].error
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_input_mapping(self):
|
|
"""Input mapping should resolve paths correctly"""
|
|
steps = [
|
|
PipelineStep(name="step1", skill="skill_a"),
|
|
PipelineStep(
|
|
name="step2",
|
|
skill="skill_b",
|
|
input_mapping={"brand": "$.input.brand"},
|
|
depends_on=["step1"],
|
|
),
|
|
]
|
|
pool = MockAgentPool({
|
|
"skill_a": MockAgent("skill_a", {"data": "a"}),
|
|
"skill_b": MockAgent("skill_b", {"data": "b"}),
|
|
})
|
|
pipeline = GEOPipeline(name="test", steps=steps, agent_pool=pool)
|
|
|
|
result = await pipeline.execute({"brand": "TestBrand"})
|
|
assert result.success
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_from_config(self):
|
|
"""Pipeline should be created from YAML config"""
|
|
config = {
|
|
"name": "geo_test",
|
|
"steps": [
|
|
{"name": "detect", "skill": "citation_detector"},
|
|
{"name": "analyze", "skill": "competitor_analyzer", "depends_on": ["detect"]},
|
|
],
|
|
}
|
|
pipeline = GEOPipeline.from_config(config)
|
|
assert pipeline.name == "geo_test"
|
|
assert len(pipeline._steps) == 2
|
|
assert pipeline._steps[1].depends_on == ["detect"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execution_groups_topological_sort(self):
|
|
"""Execution groups should follow topological order"""
|
|
steps = [
|
|
PipelineStep(name="a", skill="s1"),
|
|
PipelineStep(name="b", skill="s2", depends_on=["a"]),
|
|
PipelineStep(name="c", skill="s3", depends_on=["a"]),
|
|
PipelineStep(name="d", skill="s4", depends_on=["b", "c"]),
|
|
]
|
|
pipeline = GEOPipeline(name="test", steps=steps)
|
|
|
|
groups = pipeline._build_execution_groups()
|
|
assert len(groups) == 3
|
|
assert groups[0] == ["a"]
|
|
assert set(groups[1]) == {"b", "c"}
|
|
assert groups[2] == ["d"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_mapping_path(self):
|
|
"""Mapping path resolution"""
|
|
input_data = {"brand": "TestBrand", "platforms": ["chatgpt"]}
|
|
step_outputs = {
|
|
"detect": {"citations": 5, "records": []},
|
|
}
|
|
|
|
# $.input.brand
|
|
result = GEOPipeline._resolve_mapping_path("$.input.brand", input_data, step_outputs)
|
|
assert result == "TestBrand"
|
|
|
|
# $.steps.detect.output.citations
|
|
result = GEOPipeline._resolve_mapping_path("$.steps.detect.output.citations", input_data, step_outputs)
|
|
assert result == 5
|
|
|
|
# $.steps.detect (whole output)
|
|
result = GEOPipeline._resolve_mapping_path("$.steps.detect", input_data, step_outputs)
|
|
assert result == {"citations": 5, "records": []}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_final_output_includes_all_steps(self):
|
|
"""Final output should include all step results"""
|
|
steps = [
|
|
PipelineStep(name="step1", skill="skill_a"),
|
|
PipelineStep(name="step2", skill="skill_b", depends_on=["step1"]),
|
|
]
|
|
pool = MockAgentPool({
|
|
"skill_a": MockAgent("skill_a", {"result": "a"}),
|
|
"skill_b": MockAgent("skill_b", {"result": "b"}),
|
|
})
|
|
pipeline = GEOPipeline(name="test", steps=steps, agent_pool=pool)
|
|
|
|
result = await pipeline.execute({"query": "test"})
|
|
assert "step1" in result.final_output
|
|
assert "step2" in result.final_output
|
|
assert "input" in result.final_output
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_empty_pipeline(self):
|
|
"""Empty pipeline should succeed with no steps"""
|
|
pipeline = GEOPipeline(name="empty", steps=[])
|
|
result = await pipeline.execute({"query": "test"})
|
|
assert result.success
|
|
assert len(result.steps) == 0
|