840 lines
26 KiB
Python
840 lines
26 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""The short-term memory tests."""
|
|
import asyncio
|
|
from unittest.async_case import IsolatedAsyncioTestCase
|
|
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
|
|
|
from agentscope.memory import (
|
|
MemoryBase,
|
|
InMemoryMemory,
|
|
AsyncSQLAlchemyMemory,
|
|
RedisMemory,
|
|
)
|
|
from agentscope.message import Msg
|
|
|
|
|
|
class ShortTermMemoryTest(IsolatedAsyncioTestCase):
|
|
"""The short-term memory tests."""
|
|
|
|
memory: MemoryBase
|
|
"""The test memory instance."""
|
|
|
|
memory_session: MemoryBase
|
|
"""The test memory instance for different session."""
|
|
|
|
memory_user: MemoryBase
|
|
"""The test memory instance for different user."""
|
|
|
|
async def asyncSetUp(self) -> None:
|
|
"""Set up the memory instance for testing."""
|
|
self.msgs = [
|
|
Msg("user", "0", "user"),
|
|
Msg("user", "1", "user"),
|
|
Msg("assistant", "2", "assistant"),
|
|
Msg("system", "3", "system"),
|
|
Msg("user", "4", "user"),
|
|
Msg("assistant", "5", "assistant"),
|
|
Msg("system", "6", "system"),
|
|
Msg("user", "7", "user"),
|
|
Msg("assistant", "8", "assistant"),
|
|
Msg("system", "9", "system"),
|
|
]
|
|
for i, msg in enumerate(self.msgs):
|
|
msg.id = str(i)
|
|
|
|
async def _basic_tests(self) -> None:
|
|
"""Test the basic functionalities of the short-term memory."""
|
|
# test at the beginning
|
|
self.assertIsInstance(await self.memory.get_memory(), list)
|
|
self.assertEqual(
|
|
len(await self.memory.get_memory()),
|
|
0,
|
|
)
|
|
self.assertEqual(
|
|
await self.memory.size(),
|
|
0,
|
|
)
|
|
|
|
await self.memory.update_compressed_summary("abc")
|
|
self.assertEqual(
|
|
len(await self.memory.get_memory()),
|
|
1,
|
|
)
|
|
|
|
await self.memory.update_compressed_summary("")
|
|
self.assertEqual(
|
|
len(await self.memory.get_memory()),
|
|
0,
|
|
)
|
|
|
|
# test adding messages
|
|
await self.memory.add(self.msgs[:5])
|
|
msgs = await self.memory.get_memory()
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in range(5)],
|
|
)
|
|
|
|
# test deleting messages by id
|
|
await self.memory.delete(msg_ids=["2", "4"])
|
|
msgs = await self.memory.get_memory()
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
["0", "1", "3"],
|
|
)
|
|
self.assertEqual(
|
|
await self.memory.size(),
|
|
3,
|
|
)
|
|
|
|
# test adding more messages
|
|
await self.memory.add(self.msgs[5:])
|
|
msgs = await self.memory.get_memory()
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in [0, 1, 3, 5, 6, 7, 8, 9]],
|
|
)
|
|
|
|
# test clearing memory
|
|
await self.memory.clear()
|
|
self.assertEqual(
|
|
await self.memory.size(),
|
|
0,
|
|
)
|
|
|
|
async def _mark_tests(self) -> None:
|
|
"""Test the mark-related functionalities of the short-term memory."""
|
|
# test getting messages by nonexistent mark
|
|
await self.memory.add(self.msgs[:5])
|
|
self.assertListEqual(
|
|
[_.id for _ in await self.memory.get_memory()],
|
|
[str(_) for _ in range(5)],
|
|
)
|
|
self.assertEqual(
|
|
len(await self.memory.get_memory(mark="nonexistent")),
|
|
0,
|
|
)
|
|
|
|
# test adding marked messages
|
|
await self.memory.add(
|
|
self.msgs[5:7],
|
|
marks=["important", "todo"],
|
|
)
|
|
await self.memory.add(self.msgs[7:], marks="important")
|
|
|
|
# Test get messages by "important" mark
|
|
msgs = await self.memory.get_memory(mark="important")
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in range(5, 10)],
|
|
)
|
|
|
|
# Test get messages by "todo" mark
|
|
msgs = await self.memory.get_memory(mark="todo")
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in range(5, 7)],
|
|
)
|
|
|
|
# Test get messages excluding "todo" mark
|
|
msgs = await self.memory.get_memory(exclude_mark="todo")
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in [0, 1, 2, 3, 4, 7, 8, 9]],
|
|
)
|
|
|
|
msgs = await self.memory.get_memory(exclude_mark="important")
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in [0, 1, 2, 3, 4]],
|
|
)
|
|
|
|
# add unmarked messages
|
|
msgs = [
|
|
Msg("user", "10", "user"),
|
|
Msg("user", "11", "user"),
|
|
]
|
|
msgs[0].id = "10"
|
|
msgs[1].id = "11"
|
|
await self.memory.add(msgs)
|
|
msgs = await self.memory.get_memory()
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in range(12)],
|
|
)
|
|
|
|
# test marking messages
|
|
await self.memory.update_messages_mark(
|
|
msg_ids=["0", "1", "2"],
|
|
new_mark="review",
|
|
)
|
|
msgs = await self.memory.get_memory(mark="review")
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
["0", "1", "2"],
|
|
)
|
|
|
|
# test adding multiple marks to messages
|
|
await self.memory.update_messages_mark(
|
|
msg_ids=["6", "7", "9"],
|
|
new_mark="unread",
|
|
)
|
|
msgs = await self.memory.get_memory(mark="unread")
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in [6, 7, 9]],
|
|
)
|
|
msgs = await self.memory.get_memory(mark="important")
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in [5, 6, 7, 8, 9]],
|
|
)
|
|
|
|
# test unmarking messages
|
|
await self.memory.update_messages_mark(
|
|
msg_ids=["5", "7"],
|
|
old_mark="important",
|
|
new_mark=None,
|
|
)
|
|
self.assertListEqual(
|
|
[_.id for _ in await self.memory.get_memory(mark="important")],
|
|
[str(_) for _ in [6, 8, 9]],
|
|
)
|
|
|
|
# test updating marks
|
|
await self.memory.update_messages_mark(
|
|
msg_ids=["6", "8"],
|
|
old_mark="important",
|
|
new_mark="archived",
|
|
)
|
|
self.assertListEqual(
|
|
[_.id for _ in await self.memory.get_memory(mark="important")],
|
|
["9"],
|
|
)
|
|
self.assertListEqual(
|
|
[_.id for _ in await self.memory.get_memory(mark="archived")],
|
|
[str(_) for _ in [6, 8]],
|
|
)
|
|
|
|
# test deleting messages by mark
|
|
await self.memory.delete_by_mark("important")
|
|
msgs = await self.memory.get_memory(mark="important")
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[],
|
|
)
|
|
msgs = await self.memory.get_memory()
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 11]],
|
|
)
|
|
|
|
await self.memory.delete_by_mark(["review", "archived"])
|
|
msgs = await self.memory.get_memory()
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in [3, 4, 5, 7, 10, 11]],
|
|
)
|
|
|
|
await self.memory.clear()
|
|
msgs = await self.memory.get_memory()
|
|
self.assertEqual(
|
|
len(msgs),
|
|
0,
|
|
)
|
|
|
|
async def _multi_tenant_tests(self) -> None:
|
|
"""Test the multi-tenant functionalities of the short-term memory."""
|
|
await self.memory.add(self.msgs[:8])
|
|
msgs = await self.memory.get_memory()
|
|
self.assertEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in range(8)],
|
|
)
|
|
|
|
# Add some msgs with overlapping ids to different users' memory
|
|
await self.memory_user.add(self.msgs[3:])
|
|
self.assertEqual(
|
|
[_.id for _ in await self.memory_user.get_memory()],
|
|
[str(_) for _ in [3, 4, 5, 6, 7, 8, 9]],
|
|
)
|
|
|
|
# Mark messages
|
|
await self.memory.update_messages_mark(
|
|
new_mark="shared",
|
|
msg_ids=["5", "6", "7"],
|
|
)
|
|
|
|
# mark messages with same ids with different mark for different users
|
|
await self.memory_user.update_messages_mark(
|
|
new_mark="shared_user",
|
|
msg_ids=["6", "7", "8", "9"],
|
|
)
|
|
|
|
# Test if the marks are isolated between different users
|
|
msgs = await self.memory.get_memory(
|
|
mark="shared",
|
|
)
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in [5, 6, 7]],
|
|
)
|
|
msgs = await self.memory.get_memory(
|
|
mark="shared_user",
|
|
)
|
|
self.assertEqual(
|
|
len(msgs),
|
|
0,
|
|
)
|
|
|
|
msgs = await self.memory_user.get_memory(
|
|
mark="shared_user",
|
|
)
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in [6, 7, 8, 9]],
|
|
)
|
|
msgs = await self.memory_user.get_memory(
|
|
mark="shared",
|
|
)
|
|
self.assertEqual(
|
|
len(msgs),
|
|
0,
|
|
)
|
|
|
|
# Test delete operation is isolated between different sessions
|
|
await self.memory.delete(
|
|
msg_ids=["6", "7"],
|
|
)
|
|
msgs = await self.memory.get_memory()
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in range(6)],
|
|
)
|
|
msgs = await self.memory_user.get_memory()
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in range(3, 10)],
|
|
)
|
|
|
|
# Test delete operation by mark is isolated between different sessions
|
|
await self.memory_user.delete_by_mark("shared")
|
|
msgs = await self.memory_user.get_memory()
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in range(3, 10)],
|
|
)
|
|
msgs = await self.memory.get_memory()
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in range(6)],
|
|
)
|
|
|
|
# Clean up
|
|
await self.memory.clear()
|
|
await self.memory_user.clear()
|
|
|
|
async def _multi_session_tests(self) -> None:
|
|
"""Test the multi-session functionalities of the short-term memory."""
|
|
await self.memory.add(self.msgs[:8])
|
|
msgs = await self.memory.get_memory()
|
|
self.assertEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in range(8)],
|
|
)
|
|
|
|
# Add some msgs with overlapping ids to different session's memory
|
|
await self.memory_session.add(self.msgs[3:])
|
|
self.assertEqual(
|
|
[_.id for _ in await self.memory_session.get_memory()],
|
|
[str(_) for _ in [3, 4, 5, 6, 7, 8, 9]],
|
|
)
|
|
|
|
# Mark messages in first session
|
|
await self.memory.update_messages_mark(
|
|
new_mark="session1_mark",
|
|
msg_ids=["5", "6", "7"],
|
|
)
|
|
|
|
# mark messages with same ids with different mark for different session
|
|
await self.memory_session.update_messages_mark(
|
|
new_mark="session2_mark",
|
|
msg_ids=["6", "7", "8", "9"],
|
|
)
|
|
|
|
# Test if the marks are isolated between different sessions
|
|
msgs = await self.memory.get_memory(
|
|
mark="session1_mark",
|
|
)
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in [5, 6, 7]],
|
|
)
|
|
msgs = await self.memory.get_memory(
|
|
mark="session2_mark",
|
|
)
|
|
self.assertEqual(
|
|
len(msgs),
|
|
0,
|
|
)
|
|
|
|
msgs = await self.memory_session.get_memory(
|
|
mark="session2_mark",
|
|
)
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in [6, 7, 8, 9]],
|
|
)
|
|
msgs = await self.memory_session.get_memory(
|
|
mark="session1_mark",
|
|
)
|
|
self.assertEqual(
|
|
len(msgs),
|
|
0,
|
|
)
|
|
|
|
# Test delete operation is isolated between different sessions
|
|
await self.memory.delete(
|
|
msg_ids=["6", "7"],
|
|
)
|
|
msgs = await self.memory.get_memory()
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in range(6)],
|
|
)
|
|
msgs = await self.memory_session.get_memory()
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in range(3, 10)],
|
|
)
|
|
|
|
# Test delete operation by mark is isolated between different sessions
|
|
await self.memory_session.delete_by_mark("session1_mark")
|
|
msgs = await self.memory_session.get_memory()
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in range(3, 10)],
|
|
)
|
|
msgs = await self.memory.get_memory()
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in range(6)],
|
|
)
|
|
|
|
# Clean up
|
|
await self.memory.clear()
|
|
await self.memory_session.clear()
|
|
|
|
async def _test_add_duplicated_msgs(self) -> None:
|
|
"""Test adding duplicated messages to the memory."""
|
|
await self.memory.add(self.msgs[:8])
|
|
await self.memory.add(self.msgs[5:])
|
|
|
|
msgs = await self.memory.get_memory()
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in range(10)],
|
|
)
|
|
|
|
await self.memory.clear()
|
|
|
|
async def _test_delete_nonexistent_msg(self) -> None:
|
|
"""Test deleting nonexistent messages from the memory."""
|
|
await self.memory.add(self.msgs[:5])
|
|
await self.memory.delete(msg_ids=["nonexistent_id"])
|
|
msgs = await self.memory.get_memory()
|
|
self.assertListEqual(
|
|
[_.id for _ in msgs],
|
|
[str(_) for _ in range(5)],
|
|
)
|
|
|
|
await self.memory.clear()
|
|
|
|
async def _test_serialization(self) -> None:
|
|
"""Test the serialization and deserialization of RedisMemory."""
|
|
|
|
# Test the state dict before any updates
|
|
self.assertDictEqual(
|
|
self.memory.state_dict(),
|
|
{
|
|
"_compressed_summary": "",
|
|
},
|
|
)
|
|
|
|
# Update compressed summary and test state dict
|
|
await self.memory.update_compressed_summary("Hi there!")
|
|
state_dict = self.memory.state_dict()
|
|
|
|
# Verify the state dict content
|
|
self.assertDictEqual(
|
|
state_dict,
|
|
{
|
|
"_compressed_summary": "Hi there!",
|
|
},
|
|
)
|
|
|
|
# Clear the compressed summary and verify state dict
|
|
await self.memory.update_compressed_summary("")
|
|
self.assertDictEqual(
|
|
self.memory.state_dict(),
|
|
{
|
|
"_compressed_summary": "",
|
|
},
|
|
)
|
|
|
|
# Load the previous state dict and verify restoration
|
|
self.memory.load_state_dict(state_dict)
|
|
self.assertDictEqual(
|
|
self.memory.state_dict(),
|
|
{
|
|
"_compressed_summary": "Hi there!",
|
|
},
|
|
)
|
|
|
|
async def asyncTearDown(self) -> None:
|
|
"""Clean up after unittests"""
|
|
await self.memory.clear()
|
|
# Close the session or connection if applicable
|
|
if hasattr(self.memory, "close"):
|
|
await self.memory.close()
|
|
|
|
|
|
class InMemoryMemoryTest(ShortTermMemoryTest):
|
|
"""The in-memory short-term memory tests."""
|
|
|
|
async def asyncSetUp(self) -> None:
|
|
"""Set up the in-memory memory instance for testing."""
|
|
await super().asyncSetUp()
|
|
self.memory = InMemoryMemory()
|
|
|
|
async def test_memory(self) -> None:
|
|
"""Test the in-memory memory functionalities."""
|
|
await self._basic_tests()
|
|
await self._mark_tests()
|
|
await self._test_add_duplicated_msgs()
|
|
await self._test_delete_nonexistent_msg()
|
|
|
|
async def test_serialization(self) -> None:
|
|
"""Test the serialization and deserialization of InMemoryMemory."""
|
|
|
|
msg = Msg("user", "1", "user")
|
|
await self.memory.add(msg)
|
|
|
|
# Test the state dict before any updates
|
|
self.assertDictEqual(
|
|
self.memory.state_dict(),
|
|
{
|
|
"_compressed_summary": "",
|
|
"content": [
|
|
[
|
|
{
|
|
"id": msg.id,
|
|
"name": msg.name,
|
|
"role": msg.role,
|
|
"content": msg.content,
|
|
"metadata": msg.metadata,
|
|
"timestamp": msg.timestamp,
|
|
},
|
|
[],
|
|
],
|
|
],
|
|
},
|
|
)
|
|
|
|
# Update compressed summary and test state dict
|
|
await self.memory.update_compressed_summary("Hello World!")
|
|
state_dict = self.memory.state_dict()
|
|
|
|
# Verify the state dict content
|
|
self.assertDictEqual(
|
|
state_dict,
|
|
{
|
|
"_compressed_summary": "Hello World!",
|
|
"content": [
|
|
[
|
|
{
|
|
"id": msg.id,
|
|
"name": msg.name,
|
|
"role": msg.role,
|
|
"content": msg.content,
|
|
"metadata": msg.metadata,
|
|
"timestamp": msg.timestamp,
|
|
},
|
|
[],
|
|
],
|
|
],
|
|
},
|
|
)
|
|
|
|
# Clear the compressed summary and verify state dict
|
|
await self.memory.update_compressed_summary("")
|
|
await self.memory.clear()
|
|
self.assertDictEqual(
|
|
self.memory.state_dict(),
|
|
{
|
|
"_compressed_summary": "",
|
|
"content": [],
|
|
},
|
|
)
|
|
|
|
# Load the previous state dict and verify restoration
|
|
self.memory.load_state_dict(state_dict)
|
|
self.assertDictEqual(
|
|
self.memory.state_dict(),
|
|
{
|
|
"_compressed_summary": "Hello World!",
|
|
"content": [
|
|
[
|
|
{
|
|
"id": msg.id,
|
|
"name": msg.name,
|
|
"role": msg.role,
|
|
"content": msg.content,
|
|
"metadata": msg.metadata,
|
|
"timestamp": msg.timestamp,
|
|
},
|
|
[],
|
|
],
|
|
],
|
|
},
|
|
)
|
|
|
|
|
|
class AsyncSQLAlchemyMemoryTest(ShortTermMemoryTest):
|
|
"""The SQLAlchemy short-term memory tests."""
|
|
|
|
async def asyncSetUp(self) -> None:
|
|
"""Set up the SQLAlchemy memory instance for testing."""
|
|
await super().asyncSetUp()
|
|
self.engine = create_async_engine(
|
|
# in-memory SQLite database for testing
|
|
url="sqlite+aiosqlite:///:memory:",
|
|
)
|
|
self.memory = AsyncSQLAlchemyMemory(
|
|
session_id="session_1",
|
|
user_id="user_1",
|
|
engine_or_session=self.engine,
|
|
)
|
|
|
|
self.memory_session = AsyncSQLAlchemyMemory(
|
|
session_id="session_2",
|
|
user_id="user_1",
|
|
engine_or_session=self.engine,
|
|
)
|
|
|
|
self.memory_user = AsyncSQLAlchemyMemory(
|
|
session_id="session_2",
|
|
user_id="user_2",
|
|
engine_or_session=self.engine,
|
|
)
|
|
|
|
async def test_memory(self) -> None:
|
|
"""Test the SQLAlchemy memory functionalities."""
|
|
await self._basic_tests()
|
|
await self._test_add_duplicated_msgs()
|
|
await self._test_delete_nonexistent_msg()
|
|
await self._mark_tests()
|
|
await self._multi_tenant_tests()
|
|
await self._multi_session_tests()
|
|
await self._test_serialization()
|
|
|
|
async def test_concurrent_add(self) -> None:
|
|
"""Test that concurrent add() calls don't cause IntegrityError.
|
|
|
|
Reproduces the bug from GitHub issue #1381: when parallel_tool_calls
|
|
is True, multiple _acting coroutines call memory.add() concurrently,
|
|
causing duplicate primary key conflicts.
|
|
"""
|
|
from sqlalchemy import select as sa_select
|
|
|
|
messages = [
|
|
Msg("system", f"Tool result {i}", "system") for i in range(20)
|
|
]
|
|
|
|
# Add all messages concurrently (simulates parallel_tool_calls)
|
|
await asyncio.gather(
|
|
*(self.memory.add(msg) for msg in messages),
|
|
)
|
|
|
|
# Verify all messages were added
|
|
stored = await self.memory.get_memory()
|
|
self.assertEqual(len(stored), len(messages))
|
|
|
|
# Verify indices are unique and contiguous (the core race condition
|
|
# in _get_next_index would cause duplicate indices without the lock)
|
|
result = await self.memory.session.execute(
|
|
sa_select(self.memory.MessageTable.index)
|
|
.filter(
|
|
self.memory.MessageTable.session_id == self.memory.session_id,
|
|
)
|
|
.order_by(self.memory.MessageTable.index),
|
|
)
|
|
indices = [row[0] for row in result.fetchall()]
|
|
self.assertEqual(
|
|
len(set(indices)),
|
|
len(messages),
|
|
"Indices not unique",
|
|
)
|
|
self.assertEqual(
|
|
indices,
|
|
list(range(len(messages))),
|
|
"Indices not contiguous",
|
|
)
|
|
|
|
async def asyncTearDown(self) -> None:
|
|
"""Clean up after unittests"""
|
|
await super().asyncTearDown()
|
|
await self.engine.dispose()
|
|
|
|
|
|
class RedisMemoryTest(ShortTermMemoryTest):
|
|
"""The Redis short-term memory tests."""
|
|
|
|
memory: RedisMemory
|
|
"""The Redis memory instance."""
|
|
|
|
memory_session: RedisMemory
|
|
"""The Redis memory instance for different session."""
|
|
|
|
memory_user: RedisMemory
|
|
"""The Redis memory instance for different user."""
|
|
|
|
async def asyncSetUp(self) -> None:
|
|
"""Set up the Redis memory instance for testing."""
|
|
await super().asyncSetUp()
|
|
try:
|
|
import fakeredis.aioredis
|
|
except ImportError:
|
|
self.skipTest(
|
|
"fakeredis is not installed. Install it via "
|
|
"'pip install fakeredis' to run this test.",
|
|
)
|
|
|
|
# Use fakeredis for in-memory testing without a real Redis server
|
|
fake_redis = fakeredis.aioredis.FakeRedis(decode_responses=True)
|
|
self.memory = RedisMemory(
|
|
user_id="user_1",
|
|
session_id="session_1",
|
|
connection_pool=fake_redis.connection_pool,
|
|
)
|
|
|
|
self.memory_session = RedisMemory(
|
|
user_id="user_1",
|
|
session_id="session_2",
|
|
connection_pool=fake_redis.connection_pool,
|
|
)
|
|
|
|
self.memory_user = RedisMemory(
|
|
user_id="user_2",
|
|
session_id="session_2",
|
|
connection_pool=fake_redis.connection_pool,
|
|
)
|
|
|
|
async def test_memory(self) -> None:
|
|
"""Test the Redis memory functionalities."""
|
|
await self._basic_tests()
|
|
await self._mark_tests()
|
|
await self._test_add_duplicated_msgs()
|
|
await self._test_delete_nonexistent_msg()
|
|
await self._multi_tenant_tests()
|
|
await self._multi_session_tests()
|
|
await self._test_serialization()
|
|
|
|
async def test_ttl(self) -> None:
|
|
"""Test the TTL functionality of the Redis memory."""
|
|
# Set a short TTL for testing
|
|
self.memory.key_ttl = 2 # 2 seconds
|
|
|
|
# Add messages and verify they exist
|
|
await self.memory.add(self.msgs[:5])
|
|
msgs = await self.memory.get_memory()
|
|
self.assertEqual(
|
|
len(msgs),
|
|
5,
|
|
)
|
|
|
|
# Wait for TTL to expire
|
|
await asyncio.sleep(3)
|
|
|
|
msgs = await self.memory.get_memory()
|
|
self.assertEqual(
|
|
len(msgs),
|
|
0,
|
|
)
|
|
|
|
|
|
class RedisMemoryTestWithBytes(ShortTermMemoryTest):
|
|
"""The Redis short-term memory tests with decode_responses=False."""
|
|
|
|
memory: RedisMemory
|
|
"""The Redis memory instance."""
|
|
|
|
memory_session: RedisMemory
|
|
"""The Redis memory instance for different session."""
|
|
|
|
memory_user: RedisMemory
|
|
"""The Redis memory instance for different user."""
|
|
|
|
async def asyncSetUp(self) -> None:
|
|
"""Set up the Redis memory instance for testing."""
|
|
await super().asyncSetUp()
|
|
try:
|
|
import fakeredis.aioredis
|
|
except ImportError:
|
|
self.skipTest(
|
|
"fakeredis is not installed. Install it via "
|
|
"'pip install fakeredis' to run this test.",
|
|
)
|
|
|
|
# Use fakeredis with decode_responses=False to test bytes handling
|
|
fake_redis = fakeredis.aioredis.FakeRedis(decode_responses=False)
|
|
self.memory = RedisMemory(
|
|
user_id="user_1",
|
|
session_id="session_1",
|
|
connection_pool=fake_redis.connection_pool,
|
|
)
|
|
|
|
self.memory_session = RedisMemory(
|
|
user_id="user_1",
|
|
session_id="session_2",
|
|
connection_pool=fake_redis.connection_pool,
|
|
)
|
|
|
|
self.memory_user = RedisMemory(
|
|
user_id="user_2",
|
|
session_id="session_2",
|
|
connection_pool=fake_redis.connection_pool,
|
|
)
|
|
|
|
async def test_memory(self) -> None:
|
|
"""Test the Redis memory functionalities."""
|
|
await self._basic_tests()
|
|
await self._mark_tests()
|
|
await self._test_add_duplicated_msgs()
|
|
await self._test_delete_nonexistent_msg()
|
|
await self._multi_tenant_tests()
|
|
await self._multi_session_tests()
|
|
|
|
async def test_ttl(self) -> None:
|
|
"""Test the TTL functionality of the Redis memory."""
|
|
# Set a short TTL for testing
|
|
self.memory.key_ttl = 2 # 2 seconds
|
|
|
|
# Add messages and verify they exist
|
|
await self.memory.add(self.msgs[:5])
|
|
msgs = await self.memory.get_memory()
|
|
self.assertEqual(
|
|
len(msgs),
|
|
5,
|
|
)
|
|
|
|
# Wait for TTL to expire
|
|
await asyncio.sleep(3)
|
|
|
|
msgs = await self.memory.get_memory()
|
|
self.assertEqual(
|
|
len(msgs),
|
|
0,
|
|
)
|