tests/unit/test_session_persistence.py

"""Unit tests for the optional durable session store abstraction."""

import pytest

from agent.core.session_persistence import (
    MongoSessionStore,
    NoopSessionStore,
    _safe_message_doc,
)


@pytest.mark.asyncio
async def test_noop_store_keeps_local_cli_and_tests_db_free():
    store = NoopSessionStore()

    await store.init()
    await store.upsert_session(session_id="s1", user_id="u1", model="m")
    await store.save_snapshot(
        session_id="s1",
        user_id="u1",
        model="m",
        messages=[{"role": "user", "content": "hello"}],
    )

    assert await store.load_session("s1") is None
    assert await store.list_sessions("u1") == []
    assert await store.append_event("s1", "processing", {}) is None
    assert await store.try_increment_quota("u1", "2099-01-01", 1) is None


def test_unsafe_message_payload_is_replaced_with_marker():
    marker = _safe_message_doc({"role": "assistant", "content": object()})

    assert marker["role"] == "tool"
    assert marker["ml_intern_persistence_error"] == "message_too_large_or_invalid"


# ── mark_pro_seen ─────────────────────────────────────────────────────────


class _FakeProUsers:
    """In-memory stand-in for the ``pro_users`` collection.

    Supports just enough of the Motor API to exercise ``mark_pro_seen``:
    ``update_one`` with ``$setOnInsert`` + ``$set`` + ``upsert=True``, and
    ``find_one_and_update`` with the guarded filter the conversion check uses.
    """

    def __init__(self) -> None:
        self.docs: dict[str, dict] = {}

    async def update_one(self, filt, update, upsert=False):
        _id = filt["_id"]
        doc = self.docs.get(_id)
        if doc is None and upsert:
            doc = dict(update.get("$setOnInsert") or {})
            self.docs[_id] = doc
        if doc is None:
            return
        for k, v in (update.get("$set") or {}).items():
            doc[k] = v

    async def find_one_and_update(self, filt, update, return_document=None):
        _id = filt["_id"]
        doc = self.docs.get(_id)
        if doc is None:
            return None
        # Guard checks the conversion test uses: ever_non_pro=True AND
        # first_seen_pro_at missing.
        for k, v in filt.items():
            if k == "_id":
                continue
            if isinstance(v, dict) and "$exists" in v:
                if v["$exists"] and k not in doc:
                    return None
                if not v["$exists"] and k in doc:
                    return None
            elif doc.get(k) != v:
                return None
        for k, v in (update.get("$set") or {}).items():
            doc[k] = v
        return dict(doc)


class _FakeDB:
    def __init__(self) -> None:
        self.pro_users = _FakeProUsers()


def _store_with_fake_db() -> MongoSessionStore:
    s = MongoSessionStore.__new__(MongoSessionStore)
    s.enabled = True
    s.db = _FakeDB()
    return s


@pytest.mark.asyncio
async def test_mark_pro_seen_returns_none_when_unknown_user_starts_pro():
    """Joining as Pro shouldn't count as a conversion."""
    store = _store_with_fake_db()
    assert await store.mark_pro_seen("u-new-pro", is_pro=True) is None


@pytest.mark.asyncio
async def test_mark_pro_seen_emits_conversion_after_seeing_user_as_free():
    store = _store_with_fake_db()
    assert await store.mark_pro_seen("u1", is_pro=False) is None
    result = await store.mark_pro_seen("u1", is_pro=True)
    assert result is not None
    assert result["converted"] is True
    assert isinstance(result["first_seen_at"], str)


@pytest.mark.asyncio
async def test_mark_pro_seen_only_fires_conversion_once():
    """Re-checking a converted user must not re-emit the event."""
    store = _store_with_fake_db()
    await store.mark_pro_seen("u1", is_pro=False)
    first = await store.mark_pro_seen("u1", is_pro=True)
    assert first is not None and first["converted"] is True
    second = await store.mark_pro_seen("u1", is_pro=True)
    assert second is None


@pytest.mark.asyncio
async def test_noop_store_mark_pro_seen_returns_none():
    store = NoopSessionStore()
    assert await store.mark_pro_seen("u1", is_pro=True) is None
    assert await store.mark_pro_seen("u1", is_pro=False) is None