87 lines
3.0 KiB
Python
87 lines
3.0 KiB
Python
from pathlib import Path
|
|
|
|
from ai_orchestrator.application.ports import (
|
|
ArtifactRecord,
|
|
ModelInvocationRecord,
|
|
ToolInvocationRecord,
|
|
)
|
|
from ai_orchestrator.config import AppSettings
|
|
from ai_orchestrator.domain.events import DomainEvent
|
|
from ai_orchestrator.infrastructure.storage.factory import create_storage_bundle
|
|
|
|
|
|
def test_sqlite_paginated_event_and_invocation_queries(tmp_path: Path) -> None:
|
|
db_path = tmp_path / "queries.sqlite3"
|
|
storage = create_storage_bundle(
|
|
AppSettings(storage_backend="sqlite", database_url=f"sqlite:///{db_path.as_posix()}")
|
|
)
|
|
for index in range(5):
|
|
storage.event_store.append(
|
|
DomainEvent(event_type="test_event", task_id="task_1", payload={"index": index})
|
|
)
|
|
storage.invocation_store.save_model_invocation(
|
|
ModelInvocationRecord(
|
|
invocation_id=f"minv_{index}",
|
|
task_id="task_1",
|
|
node_id=None,
|
|
slot="weak",
|
|
provider="local",
|
|
model="mock",
|
|
status="success",
|
|
request={},
|
|
response={"message": {"content": "ok"}},
|
|
)
|
|
)
|
|
storage.invocation_store.save_tool_invocation(
|
|
ToolInvocationRecord(
|
|
invocation_id=f"tinv_{index}",
|
|
task_id="task_1",
|
|
node_id=None,
|
|
source_type="mcp",
|
|
source_id="one_c",
|
|
tool_name="one_c.run_sql",
|
|
status="success",
|
|
request={},
|
|
response={"content": {"rows": []}},
|
|
)
|
|
)
|
|
|
|
events_page = storage.event_store.list_by_task_paginated("task_1", limit=2, offset=1)
|
|
model_page = storage.invocation_store.list_model_invocations_paginated(
|
|
"task_1", limit=2, offset=2
|
|
)
|
|
tool_page = storage.invocation_store.list_tool_invocations_paginated(
|
|
"task_1", limit=3, offset=1
|
|
)
|
|
|
|
assert events_page.total == 5
|
|
assert len(events_page.items) == 2
|
|
assert model_page.total == 5
|
|
assert len(model_page.items) == 2
|
|
assert tool_page.total == 5
|
|
assert len(tool_page.items) == 3
|
|
|
|
|
|
def test_sqlite_artifact_store_persists_and_paginates(tmp_path: Path) -> None:
|
|
db_path = tmp_path / "artifacts.sqlite3"
|
|
storage = create_storage_bundle(
|
|
AppSettings(storage_backend="sqlite", database_url=f"sqlite:///{db_path.as_posix()}")
|
|
)
|
|
for index in range(4):
|
|
storage.artifact_store.save_artifact(
|
|
ArtifactRecord(
|
|
artifact_id=f"art_{index}",
|
|
task_id="task_1",
|
|
node_id=None,
|
|
artifact_type="file",
|
|
storage_uri=f"file:///artifact/{index}",
|
|
metadata={"index": index},
|
|
)
|
|
)
|
|
|
|
page = storage.artifact_store.list_by_task_paginated("task_1", limit=2, offset=1)
|
|
|
|
assert page.total == 4
|
|
assert len(page.items) == 2
|
|
assert page.items[0].storage_uri == "file:///artifact/1"
|