""" 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?