WalkXR-AI / tests / test_user_memory_store.py
test_user_memory_store.py
Raw
"""
Tests for the user-specific persistent memory store.

These tests validate that UserMemoryStore:
- Stores persistent memories in a dedicated Chroma-backed collection
- Retrieves relevant memories for a given user
- Isolates one user's memories from another user's memories
- Validates empty inputs
- Enforces the config/config_path initialization guardrail
- Automatically adds created_at timestamp metadata

To run this test, enter the following code in the terminal:
    PYTHONPATH=src pytest tests/test_user_memory_store.py -q -s
"""

from __future__ import annotations
from typing import Iterator
import pytest
from llama_index.core import Settings
from llama_index.core.embeddings import MockEmbedding
from walkxr_ai.rag.user_memory_store import UserMemoryStore


# Define a fixture saves the original embedding model and replaces it with a mock for
# testing, then restores the original after tests are complete
@pytest.fixture(autouse=True)
def configure_test_embed_model() -> Iterator[None]:
    """
    Configure a mock embedding model so memory tests run offline and deterministically.
    """
    original_embed_model = Settings._embed_model
    Settings.embed_model = MockEmbedding(embed_dim=8)
    try:
        yield
    finally:
        Settings._embed_model = original_embed_model


def test_add_and_query_single_user(tmp_path) -> None:
    """
    Test #1:
    Verify that a memory can be added for one user and later retrieved by a relevant query.

    Input:
    tmp_path: pytest fixture providing a temporary directory for the vector store to ensure test isolation.
    """
    # Store Chroma inside the temp directory called test_user_memory_store
    # When retrieving, default to the top 3 most similar results
    config = {
        "storage": {
            "persist_dir": str(tmp_path),
            "user_memory_collection_name": "test_user_memory_store",
        },
        "memory_retrieval": {
            "similarity_top_k": 3,
        },
    }

    # Create the UserMemoryStore with the test config, add a memory for user_1, and query it
    store = UserMemoryStore(config=config)
    store.add_memory(user_id="user_1", memory_text="I enjoy long walks by the lake.")
    results = store.query_memories(user_id="user_1", query="lake walks")

    # Verify that at least one relevant result comes back
    assert len(results) >= 1


def test_user_isolation(tmp_path) -> None:
    """
    Test #2:
    Verify that retrieval for one user does not return memories stored for a different user.
    """
    config = {
        "storage": {
            "persist_dir": str(tmp_path),
            "user_memory_collection_name": "test_user_memory_store",
        },
        "memory_retrieval": {
            "similarity_top_k": 3,
        },
    }

    # Add memory for two different users within the same collection, then query only for user_1
    store = UserMemoryStore(config=config)
    store.add_memory(user_id="user_1", memory_text="I love hiking")
    store.add_memory(user_id="user_2", memory_text="I hate hiking")
    results = store.query_memories(user_id="user_1", query="hiking")

    # Verify that results are returned only for user_1 and not user_2
    assert results
    assert all(result.node.metadata["user_id"] == "user_1" for result in results)
    assert all(result.node.get_text() != "I hate hiking" for result in results)


def test_empty_inputs_validation(tmp_path) -> None:
    """
    Test #3:
    Verify that empty user IDs, memory text, and query strings raise clear errors.
    """
    config = {
        "storage": {
            "persist_dir": str(tmp_path),
            "user_memory_collection_name": "test_user_memory_store",
        },
        "memory_retrieval": {
            "similarity_top_k": 3,
        },
    }

    store = UserMemoryStore(config=config)

    # Expect ValueErrors with specific messages for empty user_id, memory_text, and query inputs
    with pytest.raises(ValueError, match="user_id must be a non-empty string"):
        store.add_memory(user_id="", memory_text="Valid memory")

    with pytest.raises(ValueError, match="memory_text must be a non-empty string"):
        store.add_memory(user_id="user_1", memory_text="")

    with pytest.raises(ValueError, match="query must be a non-empty string"):
        store.query_memories(user_id="user_1", query="")


def test_config_guardrail(tmp_path) -> None:
    """
    Test #4:
    Verify that initialization fails when both config and config_path are provided.
    """
    config = {
        "storage": {
            "persist_dir": str(tmp_path),
            "user_memory_collection_name": "test_user_memory_store",
        },
        "memory_retrieval": {
            "similarity_top_k": 3,
        },
    }

    # Expect a ValueError if both config and config_path are provided to the constructor
    with pytest.raises(ValueError, match="Provide either 'config' or 'config_path'"):
        UserMemoryStore(config=config, config_path="unused.yaml")


def test_timestamp_metadata_added(tmp_path) -> None:
    """
    Test #5:
    Verify that add_memory automatically attaches created_at metadata when none is provided.
    """
    config = {
        "storage": {
            "persist_dir": str(tmp_path),
            "user_memory_collection_name": "test_user_memory_store",
        },
        "memory_retrieval": {
            "similarity_top_k": 3,
        },
    }

    # Add a memory without explicitly passing created_at metadata
    store = UserMemoryStore(config=config)
    store.add_memory(user_id="user_1", memory_text="I prefer tea in the morning.")
    results = store.query_memories(user_id="user_1", query="tea")

    # Then verify that created_at metadata was automatically added and exists in the retrieved result
    assert results
    assert "created_at" in results[0].node.metadata