web-page-backend/agent-office/app/test_db.py

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()
    ids = {a["agent_id"] for a in agents}
    # 시드된 핵심 에이전트 존재 검증 — 레지스트리 확장(insta/lotto/realestate/youtube 등)에 견고하도록
    # 고정 개수/집합이 아닌 subset으로 단언 (이전 len==2/{stock,music} 고정 단언은 stale였음).
    assert {"stock", "music"} <= ids, f"core agents missing: {ids}"
    assert len(agents) >= 2
    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")


def test_get_logs_excludes_state_messages():
    init_db()
    add_log("stock", "State: idle -> working (큐레이션 시작)")
    add_log("stock", "뉴스 12건 스크랩 완료")
    add_log("stock", "State: working -> idle ()")

    logs = get_logs("stock", limit=10)
    messages = [x["message"] for x in logs]
    assert "뉴스 12건 스크랩 완료" in messages
    assert not any(m.startswith("State: ") for m in messages)


def test_delete_old_logs_removes_beyond_retention():
    import datetime as _dt
    from app.db import delete_old_logs, _conn

    init_db()
    add_log("stock", "오래된 로그")
    # 강제로 200일 전으로 옮김
    cutoff = (_dt.datetime.utcnow() - _dt.timedelta(days=200)).isoformat()
    with _conn() as conn:
        conn.execute(
            "UPDATE agent_logs SET created_at = ? WHERE message = '오래된 로그'",
            (cutoff,),
        )

    add_log("stock", "최근 로그")
    deleted = delete_old_logs(days=90)
    assert deleted >= 1

    msgs = [x["message"] for x in get_logs("stock", limit=20)]
    assert "최근 로그" in msgs
    assert "오래된 로그" not in msgs


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