giers10/Heimgeist
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

feat(tests): Add comprehensive unit tests for agent components (registry, router, runtime, validation)

Browse Source
This commit is contained in:
Victor Giers
2026-06-15 15:24:24 +02:00
parent c1f3935cf5
commit d3ddefdc27
5 changed files with 347 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
backend/tests/__init__.py Normal file
View File

71
backend/tests/test_agent_registry.py Normal file
View File

@@ -0,0 +1,71 @@
import asyncio
import unittest
from backend.agent.registry import (
NativeToolProvider,
SchemaValidationError,
ToolDefinition,
ToolExecutionContext,
)
INPUT = {
"type": "object",
"properties": {"value": {"type": "string"}},
"required": ["value"],
"additionalProperties": False,
}
OUTPUT = {
"type": "object",
"properties": {"result": {"type": "string"}},
"required": ["result"],
"additionalProperties": False,
}
def context(registry):
return ToolExecutionContext(
run_id="run", workflow_id="workflow", node_id="node", session_id=None,
selection_mode="explicit", explicit_user_action=True,
emit=lambda *_args: asyncio.sleep(0), cancellation_event=asyncio.Event(),
db_factory=lambda: None, registry=registry,
)
class RegistryTests(unittest.IsolatedAsyncioTestCase):
async def test_duplicate_names_are_rejected(self):
async def handler(arguments, _context): return {"result": arguments["value"]}
registry = NativeToolProvider()
definition = ToolDefinition("test.echo", "echo", INPUT, OUTPUT, handler)
registry.register(definition)
with self.assertRaises(ValueError):
registry.register(definition)
async def test_arguments_and_outputs_are_validated(self):
async def handler(arguments, _context): return {"result": arguments["value"]}
registry = NativeToolProvider()
registry.register(ToolDefinition("test.echo", "echo", INPUT, OUTPUT, handler))
self.assertEqual((await registry.call_tool("test.echo", {"value": "ok"}, context(registry)))["result"], "ok")
with self.assertRaises(SchemaValidationError):
await registry.call_tool("test.echo", {"value": 1}, context(registry))
async def invalid_output(_arguments, _context): return {"result": 1}
invalid = NativeToolProvider()
invalid.register(ToolDefinition("test.invalid", "invalid", INPUT, OUTPUT, invalid_output))
with self.assertRaises(SchemaValidationError):
await invalid.call_tool("test.invalid", {"value": "x"}, context(invalid))
async def test_timeout_and_result_size_are_enforced(self):
async def slow(_arguments, _context):
await asyncio.sleep(0.05)
return {"result": "ok"}
registry = NativeToolProvider()
registry.register(ToolDefinition("test.slow", "slow", INPUT, OUTPUT, slow, timeout_seconds=0.01))
with self.assertRaises(TimeoutError):
await registry.call_tool("test.slow", {"value": "x"}, context(registry))
async def large(_arguments, _context): return {"result": "x" * 100}
limited = NativeToolProvider()
limited.register(ToolDefinition("test.large", "large", INPUT, OUTPUT, large, result_size_limit=20))
with self.assertRaises(ValueError):
await limited.call_tool("test.large", {"value": "x"}, context(limited))

60
backend/tests/test_workflow_router_builtins.py Normal file
View File

@@ -0,0 +1,60 @@
import json
from pathlib import Path
import tempfile
import unittest
from unittest.mock import AsyncMock, patch
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from backend.agent.models import WorkflowDefinition
from backend.agent.registry import NativeToolProvider
from backend.agent.router import select_workflow
from backend.agent.tools import register_native_tools
from backend.agent.validation import validate_workflow_graph
from backend.agent.workflows.builtins import BUILTIN_WORKFLOWS, seed_builtin_workflows
from backend.database import Base
from backend.ollama_client import OllamaChatResult
class RouterAndBuiltinTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.temp = tempfile.TemporaryDirectory()
engine = create_engine(f"sqlite:///{Path(self.temp.name) / 'router.db'}")
Base.metadata.create_all(engine)
self.Session = sessionmaker(bind=engine)
self.engine = engine
db = self.Session(); seed_builtin_workflows(db); db.close()
async def asyncTearDown(self):
self.engine.dispose(); self.temp.cleanup()
async def test_valid_selection_malformed_disabled_and_low_confidence_fallback(self):
db = self.Session()
web = db.query(WorkflowDefinition).filter_by(slug="web-answer").one()
with patch("backend.agent.router.chat_typed", new=AsyncMock(return_value=OllamaChatResult(content=json.dumps({"workflow_id": web.id, "confidence": 0.9, "reason": "current", "inputs": {}})))):
result = await select_workflow(db, message="latest news", recent_messages=[], attachments=[], library_slug=None, router_model=None, chat_model="model")
self.assertEqual(result["workflow_slug"], "web-answer")
with patch("backend.agent.router.chat_typed", new=AsyncMock(return_value=OllamaChatResult(content="not json"))):
self.assertEqual((await select_workflow(db, message="x", recent_messages=[], attachments=[], library_slug=None, router_model=None, chat_model="model"))["workflow_slug"], "input-output")
web.enabled = False; db.commit()
with patch("backend.agent.router.chat_typed", new=AsyncMock(return_value=OllamaChatResult(content=json.dumps({"workflow_id": web.id, "confidence": 0.9, "reason": "x", "inputs": {}})))):
self.assertEqual((await select_workflow(db, message="x", recent_messages=[], attachments=[], library_slug=None, router_model=None, chat_model="model"))["workflow_slug"], "input-output")
web.enabled = True; db.commit()
with patch("backend.agent.router.chat_typed", new=AsyncMock(return_value=OllamaChatResult(content=json.dumps({"workflow_id": web.id, "confidence": 0.2, "reason": "x", "inputs": {}})))):
self.assertEqual((await select_workflow(db, message="x", recent_messages=[], attachments=[], library_slug=None, router_model=None, chat_model="model"))["workflow_slug"], "input-output")
db.close()
async def test_missing_router_model_falls_back_to_chat_model(self):
db = self.Session()
with patch("backend.agent.router.list_model_catalog", new=AsyncMock(return_value={"chat_models": ["chat"]})), patch("backend.agent.router.chat_typed", new=AsyncMock(side_effect=RuntimeError("bad response"))) as mocked:
result = await select_workflow(db, message="x", recent_messages=[], attachments=[], library_slug=None, router_model="missing", chat_model="chat")
self.assertEqual(result["workflow_slug"], "input-output")
self.assertEqual(mocked.await_args.args[0], "chat")
db.close()
async def test_all_builtins_validate(self):
registry = NativeToolProvider(); register_native_tools(registry)
for item in BUILTIN_WORKFLOWS:
with self.subTest(item=item["slug"]):
self.assertTrue(validate_workflow_graph(item["graph"], registry).valid)

144
backend/tests/test_workflow_runtime.py Normal file
View 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")

72
backend/tests/test_workflow_validation.py Normal file
View File

@@ -0,0 +1,72 @@
import copy
import unittest
from backend.agent.registry import NativeToolProvider, ToolDefinition
from backend.agent.validation import validate_workflow_graph
async def handler(arguments, _context):
return {"content": arguments["text"]}
def registry():
value = NativeToolProvider()
value.register(ToolDefinition(
"test.echo", "echo",
{"type": "object", "properties": {"text": {"type": "string"}}, "required": ["text"], "additionalProperties": False},
{"type": "object", "properties": {"content": {"type": "string"}}, "required": ["content"], "additionalProperties": False},
handler,
))
return value
def valid_graph():
return {
"schema_version": 1,
"inputs": {"prompt": {"type": "string"}},
"outputs": {"content": {"type": "string"}},
"nodes": [
{"id": "input", "type": "input", "position": {"x": 0, "y": 0}, "config": {}},
{"id": "echo", "type": "tool", "position": {"x": 200, "y": 0}, "config": {"tool": "test.echo", "arguments": {"text": {"$ref": "input.prompt"}}}},
{"id": "output", "type": "output", "position": {"x": 400, "y": 0}, "config": {"value": {"$ref": "nodes.echo.output"}}},
],
"edges": [
{"id": "a", "source": "input", "source_handle": "output", "target": "echo", "target_handle": "input"},
{"id": "b", "source": "echo", "source_handle": "output", "target": "output", "target_handle": "input"},
],
"limits": {"maximum_nodes": 10, "maximum_tool_calls": 5, "maximum_llm_calls": 2, "maximum_runtime_seconds": 30, "maximum_result_chars": 10000, "maximum_concurrency": 2},
}
class ValidationTests(unittest.TestCase):
def setUp(self):
self.registry = registry()
def codes(self, graph):
return {item.code for item in validate_workflow_graph(graph, self.registry).errors}
def test_valid_dag(self):
self.assertTrue(validate_workflow_graph(valid_graph(), self.registry).valid)
def test_unknown_tool_missing_reference_unreachable_duplicate_and_cycle(self):
graph = valid_graph(); graph["nodes"][1]["config"]["tool"] = "missing"
self.assertIn("unknown_tool", self.codes(graph))
graph = valid_graph(); graph["nodes"][1]["config"]["arguments"]["text"] = {"$ref": "input.missing"}
self.assertIn("missing_reference", self.codes(graph))
graph = valid_graph(); graph["nodes"].append({"id": "orphan", "type": "template", "position": {"x": 0, "y": 0}, "config": {"template": "x"}})
self.assertIn("unreachable_node", self.codes(graph))
graph = valid_graph(); graph["nodes"].append(copy.deepcopy(graph["nodes"][1]))
self.assertIn("duplicate_node_id", self.codes(graph))
graph = valid_graph(); graph["edges"].append({"id": "cycle", "source": "output", "source_handle": "output", "target": "echo", "target_handle": "input"})
self.assertIn("cycle", self.codes(graph))
def test_invalid_branch_limits_and_schema(self):
graph = valid_graph()
graph["nodes"].insert(2, {"id": "condition", "type": "condition", "position": {"x": 300, "y": 0}, "config": {"value": True, "operation": "equals", "compare_to": True}})
graph["edges"][1]["target"] = "condition"
graph["edges"].append({"id": "condition-output", "source": "condition", "source_handle": "true", "target": "output", "target_handle": "input"})
self.assertIn("invalid_branch", self.codes(graph))
graph = valid_graph(); graph["limits"]["maximum_nodes"] = 999
self.assertIn("excessive_limit", self.codes(graph))
graph = valid_graph(); graph["inputs"]["prompt"] = {"type": "made-up"}
self.assertIn("malformed_schema", self.codes(graph))