init

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

@@ -0,0 +1,141 @@
+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