"""
Tests for the update_memory LangGraph node and periodic routing helper.
These tests focus on:
- Durable memory extraction node behavior
- Graceful skip/error handling
- Periodic routing every 10 turns
To run this test, enter the following code in the terminal:
PYTHONPATH=src pytest tests/test_update_memory_node.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
import walkxr_ai.core.nodes as nodes_module
from walkxr_ai.core.nodes import build_update_memory_node
from walkxr_ai.core.routing import route_after_update_conversation_history
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]:
"""
Configures 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
class FakeMemoryExtractor:
"""
Fake replacement for the real memory-extraction LLM used to test the node behavior.
This fake test double can be configured to:
- Return a specific response string (e.g. valid or invalid JSON)
- Raise an exception to simulate unexpected extractor failures
"""
def __init__(self, response: str = '{"memories": []}', error: Exception | None = None) -> None:
self.response = response
self.error = error
def extract_memories(self, system_prompt: str, prompt: str) -> str:
"""
Returns a configured response or raises a configured exception.
"""
if self.error is not None:
raise self.error
return self.response
# Define test fixtures that save the original UserMemoryStore config loading method and replace
# it with a patched version for testing, then restore the original after tests are complete
@pytest.fixture
def memory_store_config(tmp_path) -> dict:
"""
Builds an isolated config for the persistent test memory store.
This ensure test memories are written to a temporary dict, not the real ChromaDB store.
"""
return {
"storage": {
"persist_dir": str(tmp_path),
"user_memory_collection_name": "test_user_memory_store",
},
"memory_retrieval": {
"similarity_top_k": 5,
},
}
@pytest.fixture
def patch_memory_store_config(monkeypatch, memory_store_config) -> dict:
"""
Patches UserMemoryStore default config loading to point at the temp store.
"""
def fake_load_config(self, config_path: str) -> dict:
return memory_store_config
monkeypatch.setattr(nodes_module.UserMemoryStore, "_load_config", fake_load_config)
return memory_store_config
def test_update_memory_stores_extracted_memories(patch_memory_store_config, memory_store_config) -> None:
"""
Test #1:
Stores extracted memories and verifies they persist in the temporary vector store.
This test covers the full flow of the update_memory node, including:
- Calling the memory extractor
- Parsing the extractor response
- Storing extracted memories in the UserMemoryStore with correct metadata
- Retrieving stored memories to verify they were saved correctly
"""
extractor = FakeMemoryExtractor(
response=(
'{"memories": ['
'"User prefers tea in the morning.", '
'"User enjoys long reflective walks."'
"]}"
)
)
update_memory = build_update_memory_node(extractor)
state = {
"user_id": "user_1",
"conversation_history": [
{"role": "user", "content": "I always start my day with tea."},
{"role": "assistant", "content": "Tea sounds like a grounding ritual."},
{"role": "user", "content": "I also really enjoy long reflective walks."},
],
"walk_id": "walk_123",
"current_step": "conversation",
"turn_index": 9,
}
result = update_memory(state)
store = UserMemoryStore(config=memory_store_config)
results = store.query_memories(user_id="user_1", query="tea reflective walks", top_k=5)
assert result["phase"] == "memory_updated" # Is the phase correct?
assert len(results) >= 2 # Are at least 2 relevant results returned?
stored_texts = [result.node.get_text() for result in results] # Do the stored memories match the extractor response?
assert "User prefers tea in the morning." in stored_texts # Is the first memory stored correctly?
assert "User enjoys long reflective walks." in stored_texts # Is the second memory stored correctly?
assert all(result.node.metadata["user_id"] == "user_1" for result in results) # Is the user_id metadata correct?
def test_update_memory_skips_when_user_id_missing(patch_memory_store_config, memory_store_config) -> None:
"""
Test #2:
This test verifies that when user_id is missing, the node does not attempt to store memories.
"""
extractor = FakeMemoryExtractor(response='{"memories": ["User prefers tea in the morning."]}')
update_memory = build_update_memory_node(extractor)
state = {
"user_id": " ",
"conversation_history": [
{"role": "user", "content": "I like tea."},
{"role": "assistant", "content": "Tea can be comforting."},
],
"walk_id": "walk_123",
"current_step": "conversation",
"turn_index": 1,
}
result = update_memory(state)
store = UserMemoryStore(config=memory_store_config)
results = store.query_memories(user_id="user_1", query="tea", top_k=5)
assert result == {"phase": "memory_updated"} # Is the phase correct even when user_id is missing?
assert results == [] # Are no memories stored when user_id is missing?
def test_update_memory_skips_when_history_missing(patch_memory_store_config, memory_store_config) -> None:
"""
Test #3:
This test verifies that when conversation_history is missing, the node does not attempt to store memories.
"""
extractor = FakeMemoryExtractor(response='{"memories": ["User prefers tea in the morning."]}')
update_memory = build_update_memory_node(extractor)
state = {
"user_id": "user_1",
"conversation_history": [],
"walk_id": "walk_123",
"current_step": "conversation",
"turn_index": 1,
}
result = update_memory(state)
store = UserMemoryStore(config=memory_store_config)
results = store.query_memories(user_id="user_1", query="tea", top_k=5)
assert result == {"phase": "memory_updated"} # Is the phase correct even when conversation history is missing?
assert results == [] # Are no memories stored when conversation history is missing?
def test_update_memory_skips_on_invalid_json(patch_memory_store_config, memory_store_config) -> None:
"""
Test #4:
This test verifies that when the extractor response is not valid JSON, the node does
not attempt to store memories and returns a memory_updated phase without crashing.
"""
extractor = FakeMemoryExtractor(response="this response is an invalid json string that cannot be parsed")
update_memory = build_update_memory_node(extractor)
state = {
"user_id": "user_1",
"conversation_history": [
{"role": "user", "content": "I like tea."},
{"role": "assistant", "content": "Tea can be grounding."},
],
"walk_id": "walk_123",
"current_step": "conversation",
"turn_index": 1,
}
result = update_memory(state)
store = UserMemoryStore(config=memory_store_config)
results = store.query_memories(user_id="user_1", query="tea", top_k=5)
assert result == {"phase": "memory_updated"} # Is the phase correct even when JSON is invalid?
assert results == [] # Are no memories stored when extractor returns invalid JSON?
def test_update_memory_returns_error_when_extractor_raises(patch_memory_store_config) -> None:
"""
Test #5:
This test verifies that if the memory extractor raises an exception, the node catches it and returns a response with:
- Phase set to "error"
- An "error" key containing the exception message
"""
extractor = FakeMemoryExtractor(error=RuntimeError("memory extractor failed"))
update_memory = build_update_memory_node(extractor)
state = {
"user_id": "user_1",
"conversation_history": [
{"role": "user", "content": "I like tea."},
{"role": "assistant", "content": "Tea can be grounding."},
],
"walk_id": "walk_123",
"current_step": "conversation",
"turn_index": 1,
}
result = update_memory(state)
assert result["phase"] == "error" # Is the phase set to "error" when the extractor raises an exception?
assert "error" in result # Does the result include an "error" key with the exception message?
def test_route_after_update_conversation_history_triggers_on_interval() -> None:
"""
Test #6:
This test verifies that the routing helper correctly returns "update_memory" when
the completed turn index hits the defined interval (every 10 turns). For example,
if turn_index is 9 (indicating 9 completed turns), the next completed turn will be
the 10th, so the helper should return "update_memory".
"""
state = {"turn_index": 9} # Simulate that 9 turns have been completed, so the next turn will be the 10th completed turn
route = route_after_update_conversation_history(state)
assert route == "update_memory" # Does the routing helper return "update_memory" when the completed turn is on the interval?
def test_route_after_update_conversation_history_skips_off_interval() -> None:
"""
Test #7:
This test verifies that the routing helper correctly returns "end" when the completed
turn index does not hit the defined interval. For example, if turn_index is 8 (indicating
8 completed turns), the next completed turn will be the 9th, which is not on the memory
update interval, so the helper should return "end".
"""
state = {"turn_index": 8} # Simulate that 8 turns have been completed, so the next turn will be the 9th completed turn, which is not on the memory update interval
route = route_after_update_conversation_history(state)
assert route == "end" # Does the routing helper return "end" when the completed turn is not on the interval?