mirror of
https://gitee.com/wanwujie/deer-flow
synced 2026-04-05 15:10:20 +08:00
84c449cf79
* fix(checkpoint): clear in-memory store after successful persistence * test(checkpoint): add unit test for memory leak check * Update tests/unit/checkpoint/test_memory_leak.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
47 lines
1.8 KiB
Python
47 lines
1.8 KiB
Python
|
|
from unittest.mock import patch
|
|
import mongomock
|
|
import src.graph.checkpoint as checkpoint
|
|
|
|
MONGO_URL = "mongodb://admin:admin@localhost:27017/checkpointing_db?authSource=admin"
|
|
|
|
def test_memory_leak_check_memory_cleared_after_persistence():
|
|
"""
|
|
Test that InMemoryStore is cleared for a thread after successful persistence.
|
|
This prevents memory leaks for long-running processes.
|
|
"""
|
|
with patch("src.graph.checkpoint.MongoClient") as mock_mongo_client:
|
|
# Setup mongomock
|
|
mock_client = mongomock.MongoClient()
|
|
mock_mongo_client.return_value = mock_client
|
|
|
|
manager = checkpoint.ChatStreamManager(
|
|
checkpoint_saver=True,
|
|
db_uri=MONGO_URL,
|
|
)
|
|
|
|
thread_id = "leak_test_thread"
|
|
namespace = ("messages", thread_id)
|
|
|
|
# 1. Simulate streaming messages
|
|
manager.process_stream_message(thread_id, "Hello", "partial")
|
|
manager.process_stream_message(thread_id, " World", "partial")
|
|
|
|
# Verify items are in store during streaming
|
|
items = manager.store.search(namespace)
|
|
assert len(items) > 0, "Store should contain items during streaming"
|
|
|
|
# 2. Simulate end of conversation (trigger persistence)
|
|
# 'stop' should trigger _persist_complete_conversation which now includes cleanup
|
|
manager.process_stream_message(thread_id, "!", "stop")
|
|
|
|
# 3. Verify store is empty for this thread
|
|
items_after = manager.store.search(namespace)
|
|
assert len(items_after) == 0, "Memory should be cleared after successful persistence"
|
|
|
|
# Verify persistence actually happened
|
|
collection = manager.mongo_db.chat_streams
|
|
doc = collection.find_one({"thread_id": thread_id})
|
|
assert doc is not None
|
|
assert doc["messages"] == ["Hello", " World", "!"]
|