How I Slashed Chatbot Context Loss from 12% to 0 by Adding Pytest to the LLM Pipeline

It was 2 AM when the product manager’s message jolted me awake again: “Users say the context disappears in the middle of a chat. Did you change something again?” Reading the chat log where a user wrote “Wait, what did I just say?”, my blood pressure skyrocketed. This broken bug had already exploded three times, each time because the memory storage “accidentally broke”. This time I’m not fixing it — I wrote a Pytest test suite to catch the context-loss bug right at CI, so anyone who merges trash code gets a red pipeline in their face. The result: our chatbot’s context loss rate dropped from 12% to 0.

Problem breakdown

We built a memory storage module for our customer service bot. The core logic: cache the last K turns of conversation with ConversationBufferWindowMemory and periodically write the full history to PostgreSQL for persistence. When a user asks “the order number I just mentioned,” the system first checks the cache; if it’s not there, it falls back to the database.

The root cause was particularly nasty: when the cache window had fewer than K turns, the slicing logic would return incomplete history; meanwhile, database writes went through an async queue, so if the process crashed before serialization, the last few turns were permanently lost. The worst part? Developers testing manually with a few lines locally couldn’t detect it because that tiny data volume never triggered the race condition. The usual “add some logs” approach only lets you investigate after the fact — by then the users are already gone. What we needed was automated tests that could reproduce these edge cases inside CI — using Pytest to nail down the memory layer’s behavior.

Solution design

Why not use LangChain’s built-in unit tests? Because they only check that API calls succeed, not that conversational context remains semantically continuous, nor do they test concurrent window boundaries. Why not end-to-end tests? Too heavy: they require spinning up the entire chat service and take ten minutes per run — CI can’t handle that.

I chose Pytest + fixture mocking + a controlled async storage backend. The idea is simple: abstract the memory layer’s client into an interface, then inject a precision-controlled fake backend in tests that can simulate fault injection (e.g., write delays, crash-induced loss). Each test case verifies that “after multiple turns of conversation, the system can accurately recall key information from the history” — not just check the return code. Combined with parameterized tests covering various window sizes and conversation lengths, the entire CI run finishes in 36 seconds, making all context-loss bugs visible.

Core implementation

The code below defines a testable memory layer wrapper that pulls out store and retrieve interfaces, making it easy to inject fakes.

# memory_interface.py - 抽象接口，让测试可以替换真实存储
from abc import ABC, abstractmethod
from typing import List, Dict
import uuid

class Message:
    def __init__(self, role: str, content: str):
        self.role = role
        self.content = content

class MemoryBackend(ABC):
    @abstractmethod
    async def append(self, session_id: str, msg: Message):
        ...

    @abstractmethod
    async def get_history(self, session_id: str, k: int) -> List[Message]:
        ...

class ConversationMemory:
    def __init__(self, backend: MemoryBackend, window_size: int = 4):
        self.backend = backend
        self.window_size = window_size

    async def add_message(self, session_id: str, role: str, content: str):
        msg = Message(role, content)
        await self.backend.append(session_id, msg)

    async def get_context(self, session_id: str) -> List[Message]:
        return await self.backend.get_history(session_id, self.window_size)

Next is the stateful Fake Backend, which supports deliberately dropping the last N messages to simulate crash scenarios. This is the heart of the entire test suite.

# fake_backend.py - 状态可控的假存储，用于注入故障
from memory_interface import MemoryBackend, Message
from typing import List, Dict
from collections import defaultdict

class FakeMemoryBackend(MemoryBackend):
    def __init__(self, drop_last_n: int = 0):
        self._store: Dict[str, List[Message]] = defaultdict(list)
        self.drop_last_n = drop_last_n  # 模仿异步队列丢失

    async def append(self, session_id: str, msg: Message):
        self._store[session_id].append(msg)

    async def get_history(self, session_id: str, k: int) -> List[Message]:
        msgs = self._store[session_id]
        if self.drop_last_n > 0:
            # 模拟写入后崩溃，直接丢弃末尾几条
            msgs = msgs[:-self.drop_last_n]
        # 只返回最近k条，并永远保证历史是从前往后连续的一段
        return msgs[-k:] if len(msgs) >= k else msgs

Finally, the Pytest test cases cover the three most critical scenarios: incomplete window, message loss, and concurrent boundary. @pytest.mark.parametrize lets you run all window sizes at once.

# test_memory.py - 针对对话丢失的自动化测试
import pytest
from memory_interface import ConversationMemory
from fake_backend import FakeMemoryBackend

@pytest.mark.asyncio
@pytest.mark.parametrize("window_size", [2, 4, 6])