catonline_ai/vw-agentic-rag/tests/unit/test_memory.py

import pytest
from service.memory.store import InMemoryStore
from service.graph.state import TurnState, Message
from datetime import datetime, timedelta


@pytest.fixture
def memory_store():
    """Create memory store for testing"""
    return InMemoryStore(ttl_days=1)  # Short TTL for testing


@pytest.fixture
def sample_state():
    """Create sample turn state"""
    state = TurnState(session_id="test_session")
    state.messages = [
        Message(role="user", content="Hello"),
        Message(role="assistant", content="Hi there!")
    ]
    return state


def test_create_new_session(memory_store):
    """Test creating a new session"""
    state = memory_store.create_new_session("new_session")
    assert state.session_id == "new_session"
    assert len(state.messages) == 0
    
    # Verify it's stored
    retrieved = memory_store.get("new_session")
    assert retrieved is not None
    assert retrieved.session_id == "new_session"


def test_put_and_get_state(memory_store, sample_state):
    """Test storing and retrieving state"""
    memory_store.put("test_session", sample_state)
    
    retrieved = memory_store.get("test_session")
    assert retrieved is not None
    assert retrieved.session_id == "test_session"
    assert len(retrieved.messages) == 2
    assert retrieved.messages[0].content == "Hello"


def test_get_nonexistent_session(memory_store):
    """Test getting non-existent session returns None"""
    result = memory_store.get("nonexistent")
    assert result is None


def test_add_message(memory_store):
    """Test adding messages to conversation"""
    memory_store.create_new_session("test_session")
    
    message = Message(role="user", content="Test message")
    memory_store.add_message("test_session", message)
    
    state = memory_store.get("test_session")
    assert len(state.messages) == 1
    assert state.messages[0].content == "Test message"


def test_add_message_to_nonexistent_session(memory_store):
    """Test adding message creates new session if it doesn't exist"""
    message = Message(role="user", content="Test message")
    memory_store.add_message("new_session", message)
    
    state = memory_store.get("new_session")
    assert state is not None
    assert len(state.messages) == 1


def test_get_conversation_history(memory_store):
    """Test conversation history formatting"""
    memory_store.create_new_session("test_session")
    
    messages = [
        Message(role="user", content="Hello"),
        Message(role="assistant", content="Hi there!"),
        Message(role="user", content="How are you?"),
        Message(role="assistant", content="I'm doing well!")
    ]
    
    for msg in messages:
        memory_store.add_message("test_session", msg)
    
    history = memory_store.get_conversation_history("test_session")
    
    assert "User: Hello" in history
    assert "Assistant: Hi there!" in history
    assert "User: How are you?" in history
    assert "Assistant: I'm doing well!" in history


def test_get_conversation_history_empty(memory_store):
    """Test conversation history for empty session"""
    history = memory_store.get_conversation_history("nonexistent")
    assert history == ""


def test_trim_messages(memory_store):
    """Test message trimming"""
    memory_store.create_new_session("test_session")
    
    # Add many messages
    for i in range(25):
        memory_store.add_message("test_session", Message(role="user", content=f"Message {i}"))
    
    # Trim to 10 messages
    memory_store.trim("test_session", max_messages=10)
    
    state = memory_store.get("test_session")
    assert len(state.messages) <= 10


def test_trim_nonexistent_session(memory_store):
    """Test trimming non-existent session doesn't crash"""
    memory_store.trim("nonexistent", max_messages=10)
    # Should not raise an exception


def test_ttl_cleanup(memory_store):
    """Test TTL-based cleanup"""
    # Create store with very short TTL for testing
    short_ttl_store = InMemoryStore(ttl_days=0.001)  # ~1.5 minutes
    
    # Add a session
    state = TurnState(session_id="test_session")
    short_ttl_store.put("test_session", state)
    
    # Verify it exists
    assert short_ttl_store.get("test_session") is not None
    
    # Manually expire it by manipulating internal timestamp
    short_ttl_store.store["test_session"]["last_updated"] = datetime.now() - timedelta(days=1)
    
    # Try to get it - should trigger cleanup
    assert short_ttl_store.get("test_session") is None
    assert "test_session" not in short_ttl_store.store