gahusb
0613400bb7
SQLite DB layer with WAL mode, agent_config/tasks/logs/telegram_state tables, 2 seeded agents, full CRUD, and passing test suite (7/7). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
111 lines
3.7 KiB
Python
111 lines
3.7 KiB
Python
import os
|
|
import sys
|
|
import tempfile
|
|
|
|
# Override DB_PATH before importing db
|
|
_tmp = tempfile.mktemp(suffix=".db")
|
|
os.environ["AGENT_OFFICE_DB_PATH"] = _tmp
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
|
from app.db import (
|
|
init_db, get_all_agents, get_agent_config, update_agent_config,
|
|
create_task, update_task_status, approve_task, get_task, get_agent_tasks,
|
|
get_pending_approvals, add_log, get_logs,
|
|
save_telegram_callback, get_telegram_callback, mark_telegram_responded,
|
|
)
|
|
|
|
|
|
def test_init_and_seed():
|
|
init_db()
|
|
agents = get_all_agents()
|
|
assert len(agents) == 2, f"Expected 2 agents, got {len(agents)}"
|
|
ids = {a["agent_id"] for a in agents}
|
|
assert ids == {"stock", "music"}, f"Unexpected agent ids: {ids}"
|
|
print(" [PASS] test_init_and_seed")
|
|
|
|
|
|
def test_agent_config_update():
|
|
init_db()
|
|
update_agent_config("stock", custom_config={"watch": ["AAPL"]})
|
|
cfg = get_agent_config("stock")
|
|
assert cfg["custom_config"] == {"watch": ["AAPL"]}, f"Unexpected config: {cfg['custom_config']}"
|
|
print(" [PASS] test_agent_config_update")
|
|
|
|
|
|
def test_task_lifecycle():
|
|
init_db()
|
|
# Create task with approval
|
|
tid = create_task("music", "compose", {"prompt": "test"}, requires_approval=True)
|
|
task = get_task(tid)
|
|
assert task["status"] == "pending", f"Expected pending, got {task['status']}"
|
|
assert task["requires_approval"] is True
|
|
|
|
# Approve
|
|
approve_task(tid, via="telegram")
|
|
task = get_task(tid)
|
|
assert task["status"] == "approved", f"Expected approved, got {task['status']}"
|
|
assert task["approved_via"] == "telegram"
|
|
|
|
# Complete
|
|
update_task_status(tid, "succeeded", {"url": "/media/music/test.mp3"})
|
|
task = get_task(tid)
|
|
assert task["status"] == "succeeded", f"Expected succeeded, got {task['status']}"
|
|
assert task["result_data"]["url"] == "/media/music/test.mp3"
|
|
print(" [PASS] test_task_lifecycle")
|
|
|
|
|
|
def test_task_no_approval():
|
|
init_db()
|
|
tid = create_task("stock", "news_summary", {"limit": 10})
|
|
task = get_task(tid)
|
|
assert task["status"] == "working", f"Expected working, got {task['status']}"
|
|
print(" [PASS] test_task_no_approval")
|
|
|
|
|
|
def test_pending_approvals():
|
|
init_db()
|
|
create_task("music", "compose", {"prompt": "a"}, requires_approval=True)
|
|
create_task("music", "compose", {"prompt": "b"}, requires_approval=True)
|
|
create_task("stock", "news_summary", {})
|
|
pending = get_pending_approvals()
|
|
assert len(pending) == 2, f"Expected 2 pending, got {len(pending)}"
|
|
print(" [PASS] test_pending_approvals")
|
|
|
|
|
|
def test_logs():
|
|
init_db()
|
|
add_log("stock", "News fetched", "info", "task-1")
|
|
add_log("stock", "API error", "error")
|
|
logs = get_logs("stock")
|
|
assert len(logs) == 2, f"Expected 2 logs, got {len(logs)}"
|
|
assert logs[0]["level"] == "error", f"Expected error first (DESC), got {logs[0]['level']}"
|
|
print(" [PASS] test_logs")
|
|
|
|
|
|
def test_telegram_state():
|
|
init_db()
|
|
save_telegram_callback("cb-1", "task-1", "music")
|
|
cb = get_telegram_callback("cb-1")
|
|
assert cb["task_id"] == "task-1"
|
|
mark_telegram_responded("cb-1", "approve")
|
|
cb = get_telegram_callback("cb-1")
|
|
assert cb is None, f"Expected None after responded=1, got {cb}"
|
|
print(" [PASS] test_telegram_state")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_init_and_seed()
|
|
test_agent_config_update()
|
|
test_task_lifecycle()
|
|
test_task_no_approval()
|
|
test_pending_approvals()
|
|
test_logs()
|
|
test_telegram_state()
|
|
print("All DB tests passed!")
|
|
# Cleanup temp DB (best-effort; WAL mode may keep files open on Windows)
|
|
for ext in ("", "-wal", "-shm"):
|
|
try:
|
|
os.unlink(_tmp + ext)
|
|
except OSError:
|
|
pass
|