"""Tests for session memory backends."""

from yier_agents.src.memory import InMemorySessionStore, JSONSessionStore
from yier_agents.src.message import Message


def test_in_memory_session_store_round_trip():
    """Test in-memory store saves and restores messages."""
    store = InMemorySessionStore()
    messages = [
        Message(role="user", content="Hello"),
        Message(role="assistant", content="Hi there"),
    ]

    store.save("session-1", messages)

    restored = store.load("session-1")
    assert restored is not None
    assert [message.content for message in restored] == ["Hello", "Hi there"]
    assert store.get_session_count() == 1


def test_json_session_store_round_trip(tmp_path):
    """Test JSON-backed store saves and restores messages."""
    store = JSONSessionStore(tmp_path / "sessions")
    messages = [
        Message(role="system", content="System prompt"),
        Message(role="user", content="Hello"),
        Message(role="assistant", content="Hi there"),
    ]

    store.save("session-1", messages)

    restored = store.load("session-1")
    assert restored is not None
    assert [message.role for message in restored] == ["system", "user", "assistant"]
    assert store.get_session_count() == 1

    assert store.clear_session("session-1") is True
    assert store.load("session-1") is None
