118 lines
3.3 KiB
Python
118 lines
3.3 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
Unit tests for memory classes and functions
|
|
"""
|
|
|
|
import os
|
|
import unittest
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
from agentscope.message import Msg
|
|
from agentscope.memory import TemporaryMemory
|
|
|
|
|
|
class TemporaryMemoryTest(unittest.TestCase):
|
|
"""
|
|
Test cases for TemporaryMemory
|
|
"""
|
|
|
|
def setUp(self) -> None:
|
|
self.memory = TemporaryMemory()
|
|
self.file_name_1 = "tmp_mem_file1.txt"
|
|
self.file_name_2 = "tmp_mem_file2.txt"
|
|
self.msg_1 = Msg("user", "Hello", role="user")
|
|
self.msg_2 = Msg(
|
|
"agent",
|
|
"Hello! How can I help you?",
|
|
role="assistant",
|
|
)
|
|
self.msg_3 = Msg(
|
|
"user",
|
|
"Translate the following sentence",
|
|
role="assistant",
|
|
)
|
|
|
|
self.invalid = {"invalid_key": "invalid_value"}
|
|
|
|
def tearDown(self) -> None:
|
|
"""Clean up before & after tests."""
|
|
if os.path.exists(self.file_name_1):
|
|
os.remove(self.file_name_1)
|
|
if os.path.exists(self.file_name_2):
|
|
os.remove(self.file_name_2)
|
|
|
|
def test_add(self) -> None:
|
|
"""Test add different types of object"""
|
|
# add msg
|
|
self.memory.add(self.msg_1)
|
|
self.assertEqual(
|
|
self.memory.get_memory(),
|
|
[self.msg_1],
|
|
)
|
|
|
|
# add list
|
|
self.memory.add([self.msg_2, self.msg_3])
|
|
self.assertEqual(
|
|
self.memory.get_memory(),
|
|
[self.msg_1, self.msg_2, self.msg_3],
|
|
)
|
|
|
|
@patch("loguru.logger.warning")
|
|
def test_delete(self, mock_logging: MagicMock) -> None:
|
|
"""Test delete operations"""
|
|
self.memory.add([self.msg_1, self.msg_2, self.msg_3])
|
|
|
|
self.memory.delete(index=0)
|
|
self.assertEqual(
|
|
self.memory.get_memory(),
|
|
[self.msg_2, self.msg_3],
|
|
)
|
|
|
|
# test invalid
|
|
self.memory.delete(index=100)
|
|
mock_logging.assert_called_once_with(
|
|
"Skip delete operation for the invalid index [100]",
|
|
)
|
|
|
|
def test_invalid(self) -> None:
|
|
"""Test invalid operations for memory"""
|
|
# test invalid add
|
|
with self.assertRaises(Exception) as context:
|
|
self.memory.add(self.invalid)
|
|
self.assertTrue(
|
|
f"Cannot add {self.invalid} to memory" in str(context.exception),
|
|
)
|
|
|
|
def test_load_export(self) -> None:
|
|
"""
|
|
Test load and export function of TemporaryMemory
|
|
"""
|
|
memory = TemporaryMemory()
|
|
user_input = Msg(name="user", content="Hello")
|
|
agent_input = Msg(
|
|
name="agent",
|
|
content="Hello! How can I help you?",
|
|
)
|
|
memory.load([user_input, agent_input])
|
|
retrieved_mem = memory.export(to_mem=True)
|
|
self.assertEqual(
|
|
retrieved_mem,
|
|
[user_input, agent_input],
|
|
)
|
|
|
|
memory.export(file_path=self.file_name_1)
|
|
memory.clear()
|
|
self.assertEqual(
|
|
memory.get_memory(),
|
|
[],
|
|
)
|
|
memory.load(self.file_name_1)
|
|
self.assertEqual(
|
|
memory.get_memory(),
|
|
[user_input, agent_input],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|