Files
prod-end-2026/backend/tests/test_execution_service.py
T
2026-03-17 18:32:44 +03:00

694 lines
24 KiB
Python

from __future__ import annotations
import copy
from typing import Any
from uuid import uuid4
import pytest
from app.models import Action, HttpMethod
from app.models.capability import Capability
from app.models.execution import (
ExecutionRun,
ExecutionRunStatus,
ExecutionStepRun,
ExecutionStepStatus,
)
from app.models.pipeline import Pipeline, PipelineStatus
from app.services.execution_service import ExecutionService, StepExecutionError
class FakeSession:
def __init__(self, initial: dict[tuple[type[Any], Any], Any] | None = None) -> None:
self._store = initial or {}
self.step_runs_by_step: dict[int, ExecutionStepRun] = {}
self.commit_calls = 0
async def get(self, model: type[Any], key: Any) -> Any:
return self._store.get((model, key))
def add(self, obj: Any) -> None:
if isinstance(obj, ExecutionStepRun):
self.step_runs_by_step[obj.step] = obj
def add_all(self, items: list[Any]) -> None:
for item in items:
self.add(item)
async def commit(self) -> None:
self.commit_calls += 1
async def refresh(self, _obj: Any) -> None:
return None
class FakeContextStore:
def __init__(self, initial: Any = None) -> None:
self._context = initial
self.saved_contexts: list[dict[str, Any]] = []
async def load_context(self, _run_id) -> dict[str, Any]:
if isinstance(self._context, dict):
return copy.deepcopy(self._context)
return {}
async def save_context(self, _run_id, context: dict[str, Any]) -> None:
normalized = copy.deepcopy(context)
self._context = normalized
self.saved_contexts.append(normalized)
def _build_action(action_id) -> Action:
return Action(
id=action_id,
method=HttpMethod.GET,
path="/resource",
base_url="https://api.example.com",
)
def _build_capability(capability_id, action_id) -> Capability:
return Capability(
id=capability_id,
action_id=action_id,
name=f"cap_{capability_id.hex[:8]}",
)
def _build_node(step: int, capability_id, action_id, *, external_inputs: list[str] | None = None) -> dict[str, Any]:
return {
"step": step,
"name": f"Step {step}",
"external_inputs": external_inputs or [],
"endpoints": [
{
"capability_id": str(capability_id),
"action_id": str(action_id),
}
],
}
def test_topological_sort_linear_graph():
ordered = ExecutionService._topological_sort(
steps=[1, 2, 3],
edges=[
{"from_step": 1, "to_step": 2, "type": "users"},
{"from_step": 2, "to_step": 3, "type": "segments"},
],
)
assert ordered == [1, 2, 3]
def test_extract_value_from_output_by_edge_type():
output = {"users": [{"id": 1}]}
value = ExecutionService._extract_value_from_output(output, "users")
assert value == [{"id": 1}]
def test_build_request_payload_uses_path_params_and_defaults():
action = Action(
method=HttpMethod.GET,
path="/users/{user_id}",
base_url="https://api.example.com",
parameters_schema={
"type": "object",
"properties": {
"user_id": {
"type": "string",
"x-parameter-location": "path",
},
"limit": {
"type": "integer",
"x-parameter-location": "query",
"default": 10,
},
},
"required": ["user_id"],
},
)
service = ExecutionService(session=None) # type: ignore[arg-type]
payload = service._build_request_payload(
action=action,
resolved_inputs={"user_id": "abc"},
)
assert payload["url"] == "https://api.example.com/users/abc"
assert payload["query_params"] == {"limit": 10}
assert payload["missing_required"] == []
@pytest.mark.asyncio
async def test_get_action_from_node_uses_capability_action_id():
primary_action_id = uuid4()
stale_action_id = uuid4()
capability_id = uuid4()
action = _build_action(primary_action_id)
capability = _build_capability(capability_id, primary_action_id)
session = FakeSession(
{
(Capability, capability_id): capability,
(Action, primary_action_id): action,
}
)
service = ExecutionService(session=session) # type: ignore[arg-type]
node = _build_node(step=1, capability_id=capability_id, action_id=stale_action_id)
resolved_capability_id, resolved_action = await service._get_action_from_node(node)
assert resolved_capability_id == capability_id
assert resolved_action.id == primary_action_id
@pytest.mark.asyncio
async def test_get_action_from_node_raises_for_invalid_or_missing_bindings():
service = ExecutionService(session=FakeSession()) # type: ignore[arg-type]
with pytest.raises(StepExecutionError, match="valid capability_id"):
await service._get_action_from_node(
{"step": 1, "endpoints": [{"capability_id": "invalid"}]}
)
missing_capability_id = uuid4()
with pytest.raises(StepExecutionError, match=f"Capability not found: {missing_capability_id}"):
await service._get_action_from_node(
{
"step": 1,
"endpoints": [{"capability_id": str(missing_capability_id)}],
}
)
capability_id = uuid4()
capability_without_action = _build_capability(capability_id, None)
session = FakeSession({(Capability, capability_id): capability_without_action})
service = ExecutionService(session=session) # type: ignore[arg-type]
with pytest.raises(StepExecutionError, match=f"Capability does not have action_id: {capability_id}"):
await service._get_action_from_node(
{"step": 1, "endpoints": [{"capability_id": str(capability_id)}]}
)
missing_action_id = uuid4()
capability_with_missing_action = _build_capability(capability_id, missing_action_id)
session = FakeSession({(Capability, capability_id): capability_with_missing_action})
service = ExecutionService(session=session) # type: ignore[arg-type]
with pytest.raises(StepExecutionError, match=f"Action not found for capability {capability_id}: {missing_action_id}"):
await service._get_action_from_node(
{"step": 1, "endpoints": [{"capability_id": str(capability_id)}]}
)
def test_resolve_node_inputs_prefers_edge_values_over_step_outputs():
service = ExecutionService(session=None) # type: ignore[arg-type]
resolved, missing = service._resolve_node_inputs(
node={"step": 2, "external_inputs": []},
incoming_edges=[{"from_step": 1, "to_step": 2, "type": "users"}],
step_outputs={"1": {"users": [{"id": 1}]}},
edge_values={"1:2:users": [{"id": 42}]},
run_inputs={},
)
assert resolved == {"users": [{"id": 42}]}
assert missing == []
def test_resolve_node_inputs_normalizes_array_suffix_edge_types():
service = ExecutionService(session=None) # type: ignore[arg-type]
resolved, missing = service._resolve_node_inputs(
node={"step": 3, "external_inputs": []},
incoming_edges=[{"from_step": 1, "to_step": 3, "type": "users[]"}],
step_outputs={"1": {"users": [{"id": 1}]}},
edge_values={},
run_inputs={},
)
assert resolved["users[]"] == [{"id": 1}]
assert resolved["users"] == [{"id": 1}]
assert missing == []
def test_resolve_node_inputs_maps_user_hotel_pairs_to_segments():
service = ExecutionService(session=None) # type: ignore[arg-type]
segment_payload = [
{"segment_id": "seg_1", "hotel_id": "hotel_001", "user_ids": ["usr_001"]},
]
resolved, missing = service._resolve_node_inputs(
node={"step": 4, "external_inputs": []},
incoming_edges=[{"from_step": 3, "to_step": 4, "type": "user_hotel_pairs"}],
step_outputs={"3": {"segments": segment_payload}},
edge_values={},
run_inputs={},
)
assert resolved["user_hotel_pairs"] == segment_payload
assert resolved["segments"] == segment_payload
assert missing == []
def test_resolve_node_inputs_maps_empty_user_hotel_pairs_to_assignments():
service = ExecutionService(session=None) # type: ignore[arg-type]
resolved, missing = service._resolve_node_inputs(
node={"step": 5, "external_inputs": []},
incoming_edges=[{"from_step": 4, "to_step": 5, "type": "user_hotel_pairs"}],
step_outputs={"4": {"assignments": []}},
edge_values={"4:5:user_hotel_pairs": []},
run_inputs={},
)
assert resolved["user_hotel_pairs"] == []
assert resolved["assignments"] == []
assert missing == []
@pytest.mark.asyncio
async def test_execute_run_linear_pipeline_succeeds_and_persists_context():
run_id = uuid4()
pipeline_id = uuid4()
action_1_id = uuid4()
action_2_id = uuid4()
capability_1_id = uuid4()
capability_2_id = uuid4()
action_1 = _build_action(action_1_id)
action_2 = _build_action(action_2_id)
capability_1 = _build_capability(capability_1_id, action_1_id)
capability_2 = _build_capability(capability_2_id, action_2_id)
pipeline = Pipeline(
id=pipeline_id,
name="Linear pipeline",
nodes=[
_build_node(1, capability_1_id, action_1_id, external_inputs=["seed"]),
_build_node(2, capability_2_id, action_2_id),
],
edges=[{"from_step": 1, "to_step": 2, "type": "users"}],
status=PipelineStatus.READY,
)
run = ExecutionRun(
id=run_id,
pipeline_id=pipeline_id,
status=ExecutionRunStatus.QUEUED,
inputs={"seed": "abc"},
)
session = FakeSession(
{
(ExecutionRun, run_id): run,
(Pipeline, pipeline_id): pipeline,
(Capability, capability_1_id): capability_1,
(Capability, capability_2_id): capability_2,
(Action, action_1_id): action_1,
(Action, action_2_id): action_2,
}
)
context_store = FakeContextStore(initial={"step_outputs": "bad", "edge_values": []})
service = ExecutionService(session=session, context_store=context_store) # type: ignore[arg-type]
async def fake_call_action(action: Action, request_payload: dict[str, Any]):
if action.id == action_1_id:
assert request_payload["resolved_inputs"]["seed"] == "abc"
return {"status_code": 200, "body": {"users": [{"id": 1}]}}, {"users": [{"id": 1}]}
return {"status_code": 200, "body": {"ok": True}}, {"ok": True}
service._call_action = fake_call_action # type: ignore[method-assign]
await service.execute_run(run_id)
assert run.status == ExecutionRunStatus.SUCCEEDED
assert run.summary is not None
assert run.summary["total_steps"] == 2
assert run.summary["succeeded_steps"] == 2
assert run.summary["failed_steps"] == 0
assert run.summary["skipped_steps"] == 0
assert run.summary["final_output_step"] == 2
assert run.summary["final_output"] == {"ok": True}
assert session.step_runs_by_step[1].status == ExecutionStepStatus.SUCCEEDED
assert session.step_runs_by_step[2].status == ExecutionStepStatus.SUCCEEDED
assert context_store.saved_contexts[-1]["edge_values"]["1:2:users"] == [{"id": 1}]
assert context_store.saved_contexts[-1]["step_outputs"]["1"] == {"users": [{"id": 1}]}
@pytest.mark.asyncio
async def test_execute_run_is_fail_fast_and_marks_remaining_as_skipped():
run_id = uuid4()
pipeline_id = uuid4()
action_1_id = uuid4()
action_2_id = uuid4()
action_3_id = uuid4()
capability_1_id = uuid4()
capability_2_id = uuid4()
capability_3_id = uuid4()
action_1 = _build_action(action_1_id)
action_2 = _build_action(action_2_id)
action_3 = _build_action(action_3_id)
capability_1 = _build_capability(capability_1_id, action_1_id)
capability_2 = _build_capability(capability_2_id, action_2_id)
capability_3 = _build_capability(capability_3_id, action_3_id)
pipeline = Pipeline(
id=pipeline_id,
name="Fail fast pipeline",
nodes=[
_build_node(1, capability_1_id, action_1_id),
_build_node(2, capability_2_id, action_2_id),
_build_node(3, capability_3_id, action_3_id),
],
edges=[
{"from_step": 1, "to_step": 2, "type": "users"},
{"from_step": 2, "to_step": 3, "type": "segments"},
],
status=PipelineStatus.READY,
)
run = ExecutionRun(
id=run_id,
pipeline_id=pipeline_id,
status=ExecutionRunStatus.QUEUED,
inputs={},
)
session = FakeSession(
{
(ExecutionRun, run_id): run,
(Pipeline, pipeline_id): pipeline,
(Capability, capability_1_id): capability_1,
(Capability, capability_2_id): capability_2,
(Capability, capability_3_id): capability_3,
(Action, action_1_id): action_1,
(Action, action_2_id): action_2,
(Action, action_3_id): action_3,
}
)
service = ExecutionService(
session=session, # type: ignore[arg-type]
context_store=FakeContextStore(initial={"step_outputs": {}, "edge_values": {}}),
)
async def fake_call_action(action: Action, _request_payload: dict[str, Any]):
if action.id == action_2_id:
raise StepExecutionError("boom")
return {"status_code": 200}, {"users": [1]}
service._call_action = fake_call_action # type: ignore[method-assign]
await service.execute_run(run_id)
assert run.status == ExecutionRunStatus.PARTIAL_FAILED
assert run.summary is not None
assert run.summary["total_steps"] == 3
assert run.summary["succeeded_steps"] == 1
assert run.summary["failed_steps"] == 1
assert run.summary["skipped_steps"] == 1
assert run.summary["final_output_step"] == 1
assert run.summary["final_output"] == {"users": [1]}
assert session.step_runs_by_step[1].status == ExecutionStepStatus.SUCCEEDED
assert session.step_runs_by_step[2].status == ExecutionStepStatus.FAILED
assert session.step_runs_by_step[3].status == ExecutionStepStatus.SKIPPED
@pytest.mark.asyncio
async def test_execute_run_multi_endpoint_node_executes_sequential_chain():
run_id = uuid4()
pipeline_id = uuid4()
action_1_id = uuid4()
action_2_id = uuid4()
capability_1_id = uuid4()
capability_2_id = uuid4()
action_1 = Action(
id=action_1_id,
method=HttpMethod.GET,
path="/users/recent",
base_url="https://api.example.com",
)
action_2 = Action(
id=action_2_id,
method=HttpMethod.GET,
path="/segments/build",
base_url="https://api.example.com",
parameters_schema={
"type": "object",
"required": ["usersList"],
"properties": {
"usersList": {
"type": "array",
"x-parameter-location": "query",
}
},
},
)
capability_1 = _build_capability(capability_1_id, action_1_id)
capability_2 = _build_capability(capability_2_id, action_2_id)
multi_endpoint_node = {
"step": 1,
"name": "Multi endpoint node",
"external_inputs": [],
"endpoints": [
{
"capability_id": str(capability_1_id),
"action_id": str(action_1_id),
},
{
"capability_id": str(capability_2_id),
"action_id": str(action_2_id),
},
],
}
pipeline = Pipeline(
id=pipeline_id,
name="Multi endpoint chain",
nodes=[multi_endpoint_node],
edges=[],
status=PipelineStatus.READY,
)
run = ExecutionRun(
id=run_id,
pipeline_id=pipeline_id,
status=ExecutionRunStatus.QUEUED,
inputs={},
)
session = FakeSession(
{
(ExecutionRun, run_id): run,
(Pipeline, pipeline_id): pipeline,
(Capability, capability_1_id): capability_1,
(Capability, capability_2_id): capability_2,
(Action, action_1_id): action_1,
(Action, action_2_id): action_2,
}
)
service = ExecutionService(
session=session, # type: ignore[arg-type]
context_store=FakeContextStore(initial={"step_outputs": {}, "edge_values": {}}),
)
call_order: list[Any] = []
async def fake_call_action(action: Action, request_payload: dict[str, Any]):
call_order.append(action.id)
if action.id == action_1_id:
return {"status_code": 200, "body": {"users_list": [{"id": 1}]}}, {"users_list": [{"id": 1}]}
assert request_payload["resolved_inputs"]["usersList"] == [{"id": 1}]
return {"status_code": 200, "body": {"segments": [1]}}, {"segments": [1]}
service._call_action = fake_call_action # type: ignore[method-assign]
await service.execute_run(run_id)
assert run.status == ExecutionRunStatus.SUCCEEDED
assert run.summary is not None
assert run.summary["final_output"] == {"segments": [1]}
assert call_order == [action_1_id, action_2_id]
assert session.step_runs_by_step[1].capability_id == capability_1_id
assert session.step_runs_by_step[1].action_id == action_1_id
trace = session.step_runs_by_step[1].response_snapshot["endpoints_trace"] # type: ignore[index]
assert len(trace) == 2
assert trace[0]["status"] == "succeeded"
assert trace[1]["status"] == "succeeded"
@pytest.mark.asyncio
async def test_execute_run_multi_endpoint_failure_stops_pipeline():
run_id = uuid4()
pipeline_id = uuid4()
action_1_id = uuid4()
action_2_id = uuid4()
action_3_id = uuid4()
capability_1_id = uuid4()
capability_2_id = uuid4()
capability_3_id = uuid4()
action_1 = _build_action(action_1_id)
action_2 = _build_action(action_2_id)
action_3 = _build_action(action_3_id)
capability_1 = _build_capability(capability_1_id, action_1_id)
capability_2 = _build_capability(capability_2_id, action_2_id)
capability_3 = _build_capability(capability_3_id, action_3_id)
multi_endpoint_node = {
"step": 1,
"name": "Fail on second endpoint",
"external_inputs": [],
"endpoints": [
{"capability_id": str(capability_1_id), "action_id": str(action_1_id)},
{"capability_id": str(capability_2_id), "action_id": str(action_2_id)},
],
}
pipeline = Pipeline(
id=pipeline_id,
name="Failing multi-endpoint pipeline",
nodes=[
multi_endpoint_node,
_build_node(2, capability_3_id, action_3_id),
],
edges=[{"from_step": 1, "to_step": 2, "type": "segments"}],
status=PipelineStatus.READY,
)
run = ExecutionRun(
id=run_id,
pipeline_id=pipeline_id,
status=ExecutionRunStatus.QUEUED,
inputs={},
)
session = FakeSession(
{
(ExecutionRun, run_id): run,
(Pipeline, pipeline_id): pipeline,
(Capability, capability_1_id): capability_1,
(Capability, capability_2_id): capability_2,
(Capability, capability_3_id): capability_3,
(Action, action_1_id): action_1,
(Action, action_2_id): action_2,
(Action, action_3_id): action_3,
}
)
service = ExecutionService(
session=session, # type: ignore[arg-type]
context_store=FakeContextStore(initial={"step_outputs": {}, "edge_values": {}}),
)
async def fake_call_action(action: Action, _request_payload: dict[str, Any]):
if action.id == action_2_id:
raise StepExecutionError("boom")
return {"status_code": 200, "body": {"segments": [1]}}, {"segments": [1]}
service._call_action = fake_call_action # type: ignore[method-assign]
await service.execute_run(run_id)
assert run.status == ExecutionRunStatus.FAILED
assert run.summary is not None
assert run.summary["succeeded_steps"] == 0
assert run.summary["failed_steps"] == 1
assert run.summary["skipped_steps"] == 1
assert session.step_runs_by_step[1].status == ExecutionStepStatus.FAILED
assert session.step_runs_by_step[2].status == ExecutionStepStatus.SKIPPED
failed_trace = session.step_runs_by_step[1].response_snapshot["endpoints_trace"] # type: ignore[index]
assert len(failed_trace) == 2
assert failed_trace[0]["status"] == "succeeded"
assert failed_trace[1]["status"] == "failed"
@pytest.mark.asyncio
async def test_execute_run_multi_endpoint_chain_supports_composite_endpoint():
run_id = uuid4()
pipeline_id = uuid4()
action_1_id = uuid4()
atomic_capability_id = uuid4()
composite_capability_id = uuid4()
action_1 = Action(
id=action_1_id,
method=HttpMethod.GET,
path="/users/recent",
base_url="https://api.example.com",
)
atomic_capability = _build_capability(atomic_capability_id, action_1_id)
composite_capability = Capability(
id=composite_capability_id,
action_id=None,
type="COMPOSITE",
name="composite_cap",
input_schema={
"type": "object",
"required": ["users"],
"properties": {
"users": {"type": "array"},
},
},
recipe={"version": 1, "steps": [{"step": 1, "capability_id": str(atomic_capability_id), "inputs": {}}]},
)
node = {
"step": 1,
"name": "Atomic then composite",
"external_inputs": [],
"endpoints": [
{"capability_id": str(atomic_capability_id), "action_id": str(action_1_id)},
{"capability_id": str(composite_capability_id), "action_id": None},
],
}
pipeline = Pipeline(
id=pipeline_id,
name="mixed chain pipeline",
nodes=[node],
edges=[],
status=PipelineStatus.READY,
)
run = ExecutionRun(
id=run_id,
pipeline_id=pipeline_id,
status=ExecutionRunStatus.QUEUED,
inputs={},
)
session = FakeSession(
{
(ExecutionRun, run_id): run,
(Pipeline, pipeline_id): pipeline,
(Capability, atomic_capability_id): atomic_capability,
(Capability, composite_capability_id): composite_capability,
(Action, action_1_id): action_1,
}
)
service = ExecutionService(
session=session, # type: ignore[arg-type]
context_store=FakeContextStore(initial={"step_outputs": {}, "edge_values": {}}),
)
async def fake_call_action(action: Action, _request_payload: dict[str, Any]):
assert action.id == action_1_id
return {"status_code": 200, "body": {"users": [{"id": 1}]}}, {"users": [{"id": 1}]}
async def fake_execute_composite_capability(
*,
capability: Capability,
resolved_inputs: dict[str, Any],
run_inputs: dict[str, Any],
):
assert capability.id == composite_capability_id
assert resolved_inputs["users"] == [{"id": 1}]
assert run_inputs == {}
return {"capability_type": "COMPOSITE", "status_code": 200}, {"segments": [1]}
service._call_action = fake_call_action # type: ignore[method-assign]
service._execute_composite_capability = fake_execute_composite_capability # type: ignore[method-assign]
await service.execute_run(run_id)
assert run.status == ExecutionRunStatus.SUCCEEDED
assert run.summary is not None
assert run.summary["final_output"] == {"segments": [1]}
trace = session.step_runs_by_step[1].response_snapshot["endpoints_trace"] # type: ignore[index]
assert len(trace) == 2
assert trace[0]["capability_id"] == str(atomic_capability_id)
assert trace[1]["capability_id"] == str(composite_capability_id)
assert trace[1]["capability_type"] == "COMPOSITE"