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