feat(tests): Add comprehensive unit tests for agent components (registry, router, runtime, validation)
This commit is contained in:
144
backend/tests/test_workflow_runtime.py
Normal file
144
backend/tests/test_workflow_runtime.py
Normal file
@@ -0,0 +1,144 @@
|
||||
import asyncio
|
||||
import json
|
||||
from pathlib import Path
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from backend.agent.models import WorkflowDefinition, WorkflowEvent, WorkflowRevision, WorkflowRun
|
||||
from backend.agent.registry import NativeToolProvider, ToolDefinition
|
||||
from backend.agent.runtime import WorkflowRuntime
|
||||
from backend.database import Base
|
||||
from backend.models import ChatMessage, ChatSession
|
||||
|
||||
|
||||
def graph(nodes, edges, maximum_tool_calls=10):
|
||||
return {
|
||||
"schema_version": 1,
|
||||
"inputs": {"prompt": {"type": "string"}},
|
||||
"outputs": {"content": {"type": "string"}},
|
||||
"nodes": nodes,
|
||||
"edges": edges,
|
||||
"limits": {"maximum_nodes": 20, "maximum_tool_calls": maximum_tool_calls, "maximum_llm_calls": 5, "maximum_runtime_seconds": 5, "maximum_result_chars": 20000, "maximum_concurrency": 4},
|
||||
}
|
||||
|
||||
|
||||
def node(node_id, node_type, config=None):
|
||||
return {"id": node_id, "type": node_type, "position": {"x": 0, "y": 0}, "config": config or {}}
|
||||
|
||||
|
||||
def edge(source, target, handle="output"):
|
||||
return {"id": f"{source}-{target}-{handle}", "source": source, "source_handle": handle, "target": target, "target_handle": "input"}
|
||||
|
||||
|
||||
class RuntimeTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.engine = create_engine(f"sqlite:///{Path(self.temp.name) / 'test.db'}", connect_args={"check_same_thread": False})
|
||||
Base.metadata.create_all(self.engine)
|
||||
self.Session = sessionmaker(bind=self.engine)
|
||||
self.registry = NativeToolProvider()
|
||||
|
||||
async def echo(arguments, _context):
|
||||
await asyncio.sleep(arguments.get("delay", 0))
|
||||
if arguments.get("fail"):
|
||||
raise RuntimeError("requested failure")
|
||||
return {"content": arguments.get("text", ""), "sources": arguments.get("sources", [])}
|
||||
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {"text": {"type": "string"}, "delay": {"type": "number"}, "fail": {"type": "boolean"}, "sources": {"type": "array"}},
|
||||
"required": ["text"], "additionalProperties": False,
|
||||
}
|
||||
output = {"type": "object", "properties": {"content": {"type": "string"}, "sources": {"type": "array"}}, "required": ["content", "sources"], "additionalProperties": False}
|
||||
self.registry.register(ToolDefinition("test.echo", "echo", schema, output, echo))
|
||||
self.runtime = WorkflowRuntime(self.registry, self.Session)
|
||||
|
||||
async def asyncTearDown(self):
|
||||
self.engine.dispose()
|
||||
self.temp.cleanup()
|
||||
|
||||
def create_run(self, graph_value, session=True):
|
||||
db = self.Session()
|
||||
workflow = WorkflowDefinition(slug=f"w-{id(graph_value)}", name="Test", built_in=False, enabled=True)
|
||||
db.add(workflow); db.flush()
|
||||
revision = WorkflowRevision(workflow_id=workflow.id, version=1, graph_json=json.dumps(graph_value), checksum="x", created_by="test")
|
||||
db.add(revision); db.flush(); workflow.current_revision_id = revision.id
|
||||
session_id = None
|
||||
user_id = None
|
||||
if session:
|
||||
chat = ChatSession(session_id=f"s-{id(graph_value)}")
|
||||
db.add(chat); db.flush()
|
||||
user = ChatMessage(session_pk=chat.id, role="user", content="hello")
|
||||
db.add(user); db.flush(); session_id = chat.session_id; user_id = user.message_id
|
||||
run = WorkflowRun(
|
||||
workflow_id=workflow.id, workflow_revision_id=revision.id, session_id=session_id, user_message_id=user_id,
|
||||
selection_mode="explicit", inputs_json=json.dumps({"input": {"prompt": "hello"}, "run": {"messages": [{"role": "user", "content": "hello"}]}}),
|
||||
)
|
||||
db.add(run); db.commit(); run_id = run.id; db.close()
|
||||
return run_id
|
||||
|
||||
async def execute(self, run_id):
|
||||
await self.runtime.execute(run_id)
|
||||
db = self.Session(); run = db.query(WorkflowRun).filter_by(id=run_id).first(); db.expunge(run); db.close(); return run
|
||||
|
||||
async def test_basic_persistence_and_source_propagation(self):
|
||||
value = graph(
|
||||
[node("input", "input"), node("echo", "tool", {"tool": "test.echo", "arguments": {"text": {"$ref": "input.prompt"}, "sources": [{"url": "https://example.com"}]}}), node("output", "output", {"value": {"$ref": "nodes.echo.output"}})],
|
||||
[edge("input", "echo"), edge("echo", "output")],
|
||||
)
|
||||
run = await self.execute(self.create_run(value))
|
||||
self.assertEqual(run.status, "completed")
|
||||
db = self.Session()
|
||||
assistant = db.query(ChatMessage).filter(ChatMessage.role == "assistant").one()
|
||||
self.assertEqual(assistant.content, "hello")
|
||||
self.assertEqual(json.loads(assistant.sources_json)[0]["url"], "https://example.com")
|
||||
self.assertTrue(db.query(WorkflowEvent).filter(WorkflowEvent.event_type == "source_found").count())
|
||||
db.close()
|
||||
|
||||
async def test_concurrent_nodes_and_condition_branch(self):
|
||||
value = graph(
|
||||
[
|
||||
node("input", "input"),
|
||||
node("a", "tool", {"tool": "test.echo", "arguments": {"text": "a", "delay": 0.04}}),
|
||||
node("b", "tool", {"tool": "test.echo", "arguments": {"text": "b", "delay": 0.04}}),
|
||||
node("condition", "condition", {"value": True, "operation": "equals", "compare_to": True}),
|
||||
node("yes", "template", {"template": "yes"}),
|
||||
node("no", "template", {"template": "no"}),
|
||||
node("output", "output", {"value": {"content": {"$ref": "nodes.yes.output"}, "sources": []}}),
|
||||
],
|
||||
[edge("input", "a"), edge("input", "b"), edge("a", "condition"), edge("b", "condition"), edge("condition", "yes", "true"), edge("condition", "no", "false"), edge("yes", "output")],
|
||||
)
|
||||
run = await self.execute(self.create_run(value))
|
||||
self.assertEqual(run.status, "completed")
|
||||
db = self.Session()
|
||||
statuses = {row.node_id: row.status for row in db.query(__import__('backend.agent.models', fromlist=['WorkflowNodeRun']).WorkflowNodeRun).all()}
|
||||
self.assertEqual(statuses["no"], "skipped")
|
||||
db.close()
|
||||
|
||||
async def test_failure_cancellation_and_tool_limit(self):
|
||||
failing = graph(
|
||||
[node("input", "input"), node("bad", "tool", {"tool": "test.echo", "arguments": {"text": "x", "fail": True}}), node("output", "output", {"value": {"$ref": "nodes.bad.output"}})],
|
||||
[edge("input", "bad"), edge("bad", "output")],
|
||||
)
|
||||
self.assertEqual((await self.execute(self.create_run(failing))).status, "failed")
|
||||
|
||||
limited = graph(
|
||||
[node("input", "input"), node("a", "tool", {"tool": "test.echo", "arguments": {"text": "a"}}), node("b", "tool", {"tool": "test.echo", "arguments": {"text": "b"}}), node("output", "output", {"value": {"$ref": "nodes.b.output"}})],
|
||||
[edge("input", "a"), edge("a", "b"), edge("b", "output")], maximum_tool_calls=1,
|
||||
)
|
||||
self.assertEqual((await self.execute(self.create_run(limited))).status, "failed")
|
||||
|
||||
slow = graph(
|
||||
[node("input", "input"), node("slow", "tool", {"tool": "test.echo", "arguments": {"text": "x", "delay": 0.3}}), node("output", "output", {"value": {"$ref": "nodes.slow.output"}})],
|
||||
[edge("input", "slow"), edge("slow", "output")],
|
||||
)
|
||||
run_id = self.create_run(slow)
|
||||
task = asyncio.create_task(self.runtime.execute(run_id))
|
||||
await asyncio.sleep(0.03)
|
||||
self.runtime.cancel(run_id)
|
||||
await task
|
||||
db = self.Session(); status = db.query(WorkflowRun).filter_by(id=run_id).one().status; db.close()
|
||||
self.assertEqual(status, "cancelled")
|
||||
Reference in New Issue
Block a user