data-query-agent/tests/test_state.py

import time

from data_query_agent.state import StateStore


def test_conversations_are_isolated_and_resettable(tmp_path) -> None:
    store = StateStore(tmp_path / "state.sqlite3", idle_hours=24)
    first = store.get_conversation("group:c1:u1")
    second = store.get_conversation("group:c2:u1")
    assert first.session_id != second.session_id

    store.set_thread(first.key, "thread-1")
    assert store.get_conversation(first.key).thread_id == "thread-1"
    cleared = store.reset_conversation(first.key, new_session=False)
    assert cleared.session_id == first.session_id
    assert cleared.thread_id is None
    renewed = store.reset_conversation(first.key, new_session=True)
    assert renewed.session_id != first.session_id


def test_message_claim_is_idempotent(tmp_path) -> None:
    store = StateStore(tmp_path / "state.sqlite3")
    assert store.claim_message("m1", "p2p:c1:u1") is True
    assert store.claim_message("m1", "p2p:c1:u1") is False


def test_active_run_is_failed_without_overwriting_completed_run(tmp_path) -> None:
    store = StateStore(tmp_path / "state.sqlite3")
    active_dir = tmp_path / "active"
    done_dir = tmp_path / "done"
    store.create_run("r1", "p2p:c1:u1", "m1", active_dir)
    store.create_run("r2", "p2p:c1:u1", "m1", done_dir)
    store.update_run("r2", "completed")
    store.fail_active_runs("m1", "safe error")
    rows = store._conn.execute("SELECT run_id, status FROM runs ORDER BY run_id").fetchall()
    assert [(row["run_id"], row["status"]) for row in rows] == [
        ("r1", "failed"),
        ("r2", "completed"),
    ]


def test_idle_conversation_starts_new_thread(tmp_path) -> None:
    store = StateStore(tmp_path / "state.sqlite3", idle_hours=1)
    conversation = store.get_conversation("p2p:c1:u1")
    store.set_thread(conversation.key, "old-thread")
    with store._conn:
        store._conn.execute(
            "UPDATE conversations SET last_active_at=? WHERE conversation_key=?",
            (time.time() - 7200, conversation.key),
        )
    expired = store.get_conversation(conversation.key)
    assert expired.thread_id is None
    assert expired.session_id != conversation.session_id