133 lines
5.6 KiB
Python
133 lines
5.6 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from copy import deepcopy
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from agent_framework.workflows import FileWorkflowRepository, WorkflowActionRegistry, WorkflowRuntime
|
|
from agent_framework.workflows.models import WorkflowDefinition
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
CASES = ROOT / "tests" / "migration" / "data" / "workflow_yaml_cases"
|
|
|
|
|
|
def _load_cases():
|
|
return [json.loads(p.read_text(encoding="utf-8")) for p in sorted(CASES.glob("*.json"))]
|
|
|
|
|
|
def _contains_subset(expected: Any, observed: Any) -> bool:
|
|
if isinstance(expected, dict):
|
|
return isinstance(observed, dict) and all(
|
|
key in observed and _contains_subset(value, observed[key]) for key, value in expected.items()
|
|
)
|
|
if isinstance(expected, list):
|
|
return isinstance(observed, list) and len(expected) <= len(observed) and all(
|
|
_contains_subset(value, observed[index]) for index, value in enumerate(expected)
|
|
)
|
|
return expected == observed
|
|
|
|
|
|
class _DefinitionRepository:
|
|
def __init__(self, definition):
|
|
self.definition = definition
|
|
|
|
def get_active(self, name):
|
|
assert name == self.definition.name
|
|
return self.definition
|
|
|
|
def get_version(self, name, version):
|
|
assert (name, version) == (self.definition.name, self.definition.version)
|
|
return self.definition
|
|
|
|
|
|
def _normalized_expected_nodes(case):
|
|
nodes = list((case.get("expect") or {}).get("trace_nodes") or [])
|
|
# Runtime antigo repetia a action imediatamente anterior ao pause. O runtime
|
|
# novo deliberadamente não a repete para impedir side effects duplicados.
|
|
normalized = []
|
|
for node in nodes:
|
|
if normalized and normalized[-1] == node:
|
|
continue
|
|
normalized.append(node)
|
|
return normalized
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("case", _load_cases(), ids=lambda c: c["id"])
|
|
async def test_original_workflow_case_on_framework_runtime(case):
|
|
original = FileWorkflowRepository(ROOT / "workflows").get_version(case["workflow"], int(case["version"]))
|
|
raw = original.model_dump(by_alias=True)
|
|
registry = WorkflowActionRegistry()
|
|
calls = []
|
|
|
|
# No teste cada node recebe uma action exclusiva. Isso permite reproduzir
|
|
# exatamente os outputs configurados no caso JSON, mesmo quando dois nodes
|
|
# compartilham a mesma action de produção.
|
|
for node in raw["nodes"]:
|
|
node_id = node["id"]
|
|
real_action = node["action"]
|
|
test_action = f"__case__{node_id}"
|
|
node["action"] = test_action
|
|
|
|
async def handler(params, state, *, _node=node_id, _action=real_action):
|
|
cfgs = dict(case.get("actions") or {})
|
|
configured = cfgs.get(_node) if isinstance(cfgs.get(_node), dict) else cfgs.get(_action)
|
|
configured = configured if isinstance(configured, dict) else {}
|
|
calls.append({"node_id": _node, "action": _action, "params": deepcopy(params)})
|
|
if configured.get("success") is False:
|
|
raise RuntimeError(str(configured.get("error") or "falha configurada"))
|
|
return deepcopy(configured.get("output") or {"success": True})
|
|
|
|
registry.register(test_action, handler)
|
|
|
|
definition = WorkflowDefinition.model_validate(raw)
|
|
runtime = WorkflowRuntime(_DefinitionRepository(definition), actions=registry, checkpointer=None, allow_deterministic_fallback=True)
|
|
first = await runtime.arun(definition.name, deepcopy(case.get("input") or {}), version=definition.version)
|
|
|
|
start_expect = case.get("expect_start")
|
|
if isinstance(start_expect, dict):
|
|
if start_expect.get("status") == "WAITING_INPUT":
|
|
assert first.status == "PAUSED"
|
|
if "paused_at" in start_expect:
|
|
assert first.pause.get("node") == start_expect["paused_at"]
|
|
expected_input = first.pause.get("expected_input") or {}
|
|
if "expected_input_key" in start_expect:
|
|
assert expected_input.get("key") == start_expect["expected_input_key"]
|
|
if "allowed_values" in start_expect:
|
|
assert expected_input.get("allowed_values") == start_expect["allowed_values"]
|
|
|
|
final = first
|
|
resume = case.get("resume")
|
|
if isinstance(resume, dict):
|
|
final = await runtime.aresume(definition.name, first.execution_id, deepcopy(resume), version=definition.version)
|
|
|
|
expect = dict(case.get("expect") or {})
|
|
assert final.status == expect.get("status")
|
|
if "last_node" in expect:
|
|
assert final.state.get("current_node") == expect["last_node"]
|
|
|
|
# Data do runtime antigo era a união dos outputs. ``vars`` representa essa
|
|
# informação no runtime novo e é a fonte para referências $.vars.* do YAML.
|
|
flattened = {}
|
|
for value in (final.state.get("vars") or {}).values():
|
|
if isinstance(value, dict):
|
|
flattened.update(value)
|
|
if "data_subset" in expect:
|
|
assert _contains_subset(expect["data_subset"], flattened)
|
|
|
|
observed_nodes = [item.get("node") for item in final.trace if item.get("action") != "pause_resume"]
|
|
if "trace_nodes" in expect:
|
|
active_node_ids = {node.id for node in definition.nodes}
|
|
expected_nodes = [node for node in _normalized_expected_nodes(case) if node in active_node_ids]
|
|
assert observed_nodes == expected_nodes
|
|
|
|
for expected in expect.get("calls", []):
|
|
found = [c for c in calls if c["node_id"] == expected.get("node_id") and ("action" not in expected or c["action"] == expected["action"])]
|
|
assert found, expected
|
|
subset = expected.get("params_subset")
|
|
if subset is not None:
|
|
assert any(_contains_subset(subset, c["params"]) for c in found)
|