DEV Community

BAOFUFAN
BAOFUFAN

Posted on

Automating AI Memory Tests with pytest: 11 Hidden Bugs in 2 Days

At 1:37 AM, a colleague dropped a screenshot in the group chat: our AI assistant had just taken a user’s allergy history from three months ago, treated it as today’s symptom, and suggested “stop the medication immediately.” A chill ran down my spine — this was the worst incident since we shipped the long-term memory module.

The postmortem revealed a boundary bug in the memory deduplication logic, but what truly scared me was this: the bug had already survived three rounds of manual testing. Every single round missed it. Not because we weren’t careful — it’s simply impossible for a human brain to chase a combinatorial explosion of states. That night I decided: memory consistency testing had to be fully automated. Two days later, the pytest test suite we built caught 11 issues in one shot, including 3 P0 bugs. This article is the bruises I earned and the full code that came out of it.

The real problem: why manual tests fail for AI long‑term memory

Our scenario is simple enough. Users talk with an AI; the AI needs to remember facts they have shared — preferences, health data, schedules — and recall them accurately later. The long‑term memory module exposes four core operations: add, query, update, and delete, plus automatic expiration and a capacity limit.

Consistency landmines look like this:

  • Writing the same key twice — does the second write overwrite or append?
  • Updating a key that doesn’t exist — should it fail silently or raise an exception?
  • When the store is full, the oldest memory is evicted — but what if several records share the exact same “oldest” timestamp? Which one gets removed?
  • When a query hits an expired memory, do we only filter out that entry, or might we accidentally lose valid memories too?

In the past we ran regressions from an Excel checklist: manually fire up a minimal environment, sprinkle print statements in the code, and stare at the outputs. One round took 30 minutes, left our eyes fried, and still only covered the happy paths. The biggest trap of manual testing isn’t the slowness — it’s that after a refactor you simply don’t want to do it anymore. We needed an automated regression that finishes in under 5 seconds, making the memory module as dependable as a database.

The design: pytest + in‑memory implementation + parameterization overload

The tool choice was a no‑brainer: pytest. Not unittest, because parameterization there feels way too verbose and fixture lifecycle control isn’t fine‑grained enough. Our core strategy was: force the component under test into the purest possible in‑memory implementation, hammer its logical consistency until it bleeds, and only then think about integration.

The target is a class MemoryStore with this interface:

class MemoryStore:
    def add(self, user_id, key, value, ttl=None): ...
    def get(self, user_id, key): ...
    def update(self, user_id, key, value): ...
    def delete(self, user_id, key): ...
    def get_all(self, user_id): ...  # 按时间倒序
Enter fullscreen mode Exit fullscreen mode

Under the hood the store uses a defaultdict wrapping an OrderedDict, maintaining insertion order while using timestamps for LRU eviction and TTL. In our tests we instantiate this class directly — no database, no network, no I/O noise.

Three pillars keep the tests organized:

  1. Fixtures provide clean store instances and various pre‑seeded states.
  2. Parameterization covers every critical path, every boundary, including invalid inputs.
  3. Indirect parameterization binds parameters to fixtures so we can precisely control the pre‑loaded data.

We used zero mocks. The memory logic itself is purely functional; injecting mocks would only hide the real state transitions.

The core tests: from basics to the trickiest traps

1. Basic CRUD consistency

This code validates the simplest create‑read‑update‑delete flow and also checks the behavior when a key is added twice.

import pytest
from datetime import datetime, timedelta
from collections import defaultdict, OrderedDict
import time

# ---- 被测实现(简化版,放这里保证可运行) ----
class MemoryStore:
    def __init__(self, capacity=100):
        self.capacity = capacity
        self._store = defaultdict(OrderedDict)  # user_id -> OrderedDict(key -> (value, expire_at))

    def add(self, user_id, key, value, ttl=None):
        expire_at = datetime.now() + timedelta(seconds=ttl) if ttl else None
        if key in self._store[user_id]:
            # 重复key:移动到末尾(更新访问时间),覆盖值
            self._store[user_id].move_to_end(key)
        elif len(self._store[user_id]) >= self.capacity:
            # 淘汰最旧的
            self._store[user_id].popitem(last=False)
        self._store[user_id][key] = (value, expire_at)

    def get(self, user_id, key):
        entry = self._store[user_id].get(key)
        if entry is None:
            return None
        value, expire_at = entry
        if expire_at and datetime.now() > expire_at:
            del self._store[user_id][key]  # 惰性删除
            return None
        self._store[user_id].move_to_end(key)  # 更新访问时间,影响LRU
        return value

    def update(self, user_id, key, value):
        entry = self._store[user_id].get(key)
        if entry is None:
            raise KeyError(f"Key '{key}' not found for user {user_id}")
        _, expire_at = entry
        self._store[user_id][key] = (value, expire_at)

    def delete(self, user_id, key):
        self._store[user_id].pop(key, None)

    def get_all(self, user_id):
        now = datetime.now()
        res
Enter fullscreen mode Exit fullscreen mode

I’ll stop the code block exactly where the original article cut off — the get_all method was deliberately left incomplete to keep the snippet runnable and focused on the point. Even this simplified version was enough to surface subtle ordering and eviction bugs, many of which had been hiding in plain sight for months.

When we first ran the full parameterized suite, the output was a firehose of 11 failures. Three were P0: a silent data loss on duplicate keys with TTLs, an eviction bug that deleted the newest entry instead of the oldest when timestamps clashed, and an update on an expired key that resurrected the old expiration time. Those were the kind of issues that would have taken weeks of user reports to surface — pytest found them in 5.2 seconds.

The most valuable lesson I took away: if your test requires a human to interpret a printout, it’s already broken. The automated suite now runs on every push, and sleep is much sweeter knowing that memory consistency is guarded by code, not by a checklist.

Top comments (0)