121 lines
4.1 KiB
Python
121 lines
4.1 KiB
Python
from pathlib import Path
|
|
|
|
from ai_orchestrator.application.ports import ModelInvocationRecord, ToolInvocationRecord
|
|
from ai_orchestrator.config import AppSettings
|
|
from ai_orchestrator.domain.enums import NodeType, PolicyDecisionType
|
|
from ai_orchestrator.domain.events import DomainEvent
|
|
from ai_orchestrator.domain.models import (
|
|
ConfirmationRequest,
|
|
ExecutionGraph,
|
|
ExecutionNode,
|
|
PolicyDecision,
|
|
Task,
|
|
WorkerSession,
|
|
)
|
|
from ai_orchestrator.infrastructure.storage.factory import create_storage_bundle
|
|
|
|
|
|
def test_sqlite_storage_persists_task_graph_and_events(tmp_path: Path) -> None:
|
|
db_path = tmp_path / "orchestrator.sqlite3"
|
|
settings = AppSettings(
|
|
storage_backend="sqlite",
|
|
database_url=f"sqlite:///{db_path.as_posix()}",
|
|
)
|
|
storage = create_storage_bundle(settings)
|
|
|
|
task = Task(project_id="default", goal="Persist me", inputs={"a": 1})
|
|
storage.task_repository.create(task)
|
|
node = ExecutionNode(
|
|
task_id=task.task_id,
|
|
node_type=NodeType.PLANNER,
|
|
input_data={"goal": "Persist me"},
|
|
)
|
|
graph = ExecutionGraph(task_id=task.task_id, nodes=[node])
|
|
storage.graph_repository.save(graph)
|
|
storage.event_store.append(
|
|
DomainEvent(event_type="task_created", task_id=task.task_id, payload={"goal": task.goal})
|
|
)
|
|
|
|
stored_task = storage.task_repository.get(task.task_id)
|
|
stored_graph = storage.graph_repository.get(task.task_id)
|
|
stored_events = storage.event_store.list_by_task(task.task_id)
|
|
|
|
assert stored_task is not None
|
|
assert stored_task.goal == "Persist me"
|
|
assert stored_graph is not None
|
|
assert stored_graph.nodes[0].node_type == NodeType.PLANNER
|
|
assert stored_events[0].event_type == "task_created"
|
|
|
|
|
|
def test_sqlite_storage_persists_confirmation_and_worker(tmp_path: Path) -> None:
|
|
db_path = tmp_path / "orchestrator.sqlite3"
|
|
settings = AppSettings(
|
|
storage_backend="sqlite",
|
|
database_url=f"sqlite:///{db_path.as_posix()}",
|
|
)
|
|
storage = create_storage_bundle(settings)
|
|
confirmation = ConfirmationRequest(
|
|
task_id="task_1",
|
|
node_id="node_1",
|
|
decision=PolicyDecision(
|
|
decision=PolicyDecisionType.CONFIRM,
|
|
reason="Need approval",
|
|
requires_confirmation=True,
|
|
),
|
|
)
|
|
worker = WorkerSession(
|
|
worker_id="worker_home_pc",
|
|
name="Home PC",
|
|
machine="DESKTOP-1",
|
|
os="windows",
|
|
version="0.1.0",
|
|
capabilities=["file.read"],
|
|
)
|
|
|
|
storage.confirmation_repository.create(confirmation)
|
|
storage.worker_repository.save(worker)
|
|
|
|
assert storage.confirmation_repository.get(confirmation.confirmation_id) is not None
|
|
assert storage.worker_repository.get(worker.session_id) is not None
|
|
assert len(storage.worker_repository.list_active()) == 1
|
|
|
|
|
|
def test_sqlite_storage_persists_invocation_records(tmp_path: Path) -> None:
|
|
db_path = tmp_path / "orchestrator.sqlite3"
|
|
settings = AppSettings(
|
|
storage_backend="sqlite",
|
|
database_url=f"sqlite:///{db_path.as_posix()}",
|
|
)
|
|
storage = create_storage_bundle(settings)
|
|
|
|
storage.invocation_store.save_model_invocation(
|
|
ModelInvocationRecord(
|
|
invocation_id="minv_1",
|
|
task_id="task_1",
|
|
node_id="node_1",
|
|
slot="weak",
|
|
provider="local",
|
|
model="mock",
|
|
status="success",
|
|
request={"messages": []},
|
|
response={"message": {"content": "ok"}},
|
|
usage={"input_tokens": 1},
|
|
)
|
|
)
|
|
storage.invocation_store.save_tool_invocation(
|
|
ToolInvocationRecord(
|
|
invocation_id="tinv_1",
|
|
task_id="task_1",
|
|
node_id="node_2",
|
|
source_type="mcp",
|
|
source_id="one_c",
|
|
tool_name="one_c.run_sql",
|
|
status="success",
|
|
request={"args": {"query": "select 1"}},
|
|
response={"content": {"rows": []}},
|
|
)
|
|
)
|
|
|
|
assert len(storage.invocation_store.list_model_invocations("task_1")) == 1
|
|
assert len(storage.invocation_store.list_tool_invocations("task_1")) == 1
|