At 3 a.m., our ops channel blew up: the smart customer-support bot suddenly started spouting nonsense. A user asked “When will my order arrive?” and it replied “Last night’s sunset was beautiful.” We scrambled through the logs and found that LangChain’s ConversationBufferMemory had crammed over 1,200 context entries into its buffer. The token count blew the window to pieces, leaving the LLM with fragments from three or four turns ago—both short-term and long-term memory completely scrambled. And guess what? Our previous manual regression suite never touched this scenario, because it only covered 5 conversation flows end-to-end.
If you’re using Memory to make your LLM “remember conversations,” a blow-up like this is just a matter of time. Testing memory by hand is like crossing the street blindfolded—you might get lucky, or you might cause a production incident. We needed something automated, repeatable, and capable of covering extreme edge cases: treat the LLM’s short/long-term memory store as a stateful distributed system, then pin down its correctness with pytest assertions. Once we did this, our team’s defect discovery rate shot from an average of 2 per month (manual era) to 7 on the very first automated run, and the regression cycle shrank from 4 hours to 3 minutes. Below I’ll unpack the entire approach—you can grab the code and run it right away.
The problem isn’t just "forgetting"—it’s "mixing things up"
In LLM apps, the Memory component is responsible for preserving and passing context across multiple calls. Three typical implementations exist: raw buffer memory (e.g., ConversationBufferMemory that keeps the full history), summary memory (ConversationSummaryMemory, which distills the conversation every N turns), and window + token-count memory (ConversationTokenBufferMemory). Looks ideal, right? But as soon as you go to production, you hit every pothole:
- Short-memory overwrite: After many turns, old context that should have been evicted still lingers in some key, causing the LLM to mix up conversations.
- Long-memory misalignment: When the user switches topics, the summary doesn’t update, so long-term memory grafts key facts from one topic onto another.
- Concurrency safety: When two messages arrive in quick succession under the same session ID, the underlying dictionary operations inside Memory aren’t atomic, so the context gets corrupted.
-
Inaccurate token counting: The token count computed by
tiktokendrifts from what the model actually consumes; the window looks half-full but the API has already rejected the request.
How painful is testing this by hand? You write a small script, simulate 5–10 conversation turns, inspect the output with your own eyes, then manually tweak a few messages to see whether memory updates correctly. One round takes half an hour, only covers happy paths, and forget about boundary values. Even worse, when you swap the underlying model (say from GPT-3.5 to GPT-4), the token-counting rules change, and all your manual cases are effectively invalidated. The root cause: we never locked down memory’s behavioral contract with automated assertions.
Design: test memory like you’d test a database
We don’t “test the LLM”; we test the logic layer of memory storage. That means we mock the LLM calls and focus exclusively on whether the Memory object, after repeated save_context and load_memory_variables calls, maintains correct internal state and returns context strings that satisfy a consistency contract.
The tech stack is simple: pytest + LangChain’s memory suite + unittest.mock. Not Jupyter Notebooks—those aren’t regression-friendly. Not unittest alone—pytest’s parametrize and fixture let the same test case auto-expand across 4 memory implementations and 3 token counters, while injecting parameterized fake conversation histories. Doing that by hand is impossible.
Architecture idea: define an abstract test base class BaseMemoryTests that aggregates 10+ orthogonal test methods. Each method verifies one rigid property of memory, e.g.: does the stored length exceed the limit after multiple writes? Are old messages truncated as expected? Does summary update preserve core entities? Is the state truly clean after a reset? Then let each concrete memory implementation inherit this base and inject its own initialization logic. Finally, use pytest_generate_tests to dynamically create parameterized test cases, and let the full assertion suite run in the middle of the night.
There’s one more power move: replace the LLM summarization model with a deterministic function. For example, when testing ConversationSummaryMemory, we don’t call OpenAI; we inject an LLM stub that returns a fixed string. That makes the summary content 100% predictable, allowing us to write hard assertions. Otherwise, anything based on “semantic similarity” makes test stability a nightmare.
Core implementation: it works the moment you run it
Step 1: abstract test base class — lock down memory’s behavioral contract
This code enforces “basic rules that every memory implementation must follow.” We define BaseTestMemory; any new memory backend just inherits it and implements make_memory, and it automatically gets this whole test matrix.
# test_memory_base.py
import pytest
from langchain.schema import HumanMessage, AIMessage
from langchain.memory.chat_memory import BaseChatMemory
class BaseTestMemory:
"""所有memory实现的公共测试套件"""
memory: BaseChatMemory = None # 子类实现 make_memory 注入
def _add_turns(self, turns: int):
"""注入标准多轮对话,模拟真实交互"""
for i in range(turns):
self.memory.save_context(
{"input": f"用户消息{i}"},
{"output": f"助手回答{i}"}
)
def test_single_turn_memory(self):
"""单轮对话后,历史中应包含对应的Human和AI消息"""
self._add_turns(1)
variables = self.memory.load_memory_variables({})
history_text = variables.get("history", "")
assert "用户消息0" in history_text
assert "助手回答0" in history_text
def test_overflow_oldest_eviction(self):
"""窗口记忆溢出时,应优先丢弃最早的消息"""
# 假设 make_memory 创建了一个窗口为 4 条消息的记忆(2轮)
self._add_turns(3) # 存3轮,共6条消息
variables = self.memory.load_memory_variables({})
history = variables.get("history", "")
# 验证最旧的第一轮已被踢出,但最新的第三轮还在
assert "用户消息0" not in history
Top comments (0)