250 lines
9.8 KiB
Python
250 lines
9.8 KiB
Python
"""Tests for MemoryFile + MemoryStore — 记忆文件读写与多文件管理 (U1+U2)."""
|
|
|
|
import tempfile
|
|
from datetime import datetime, timedelta, timezone
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from agentkit.memory.profile import MemoryFile, MemoryStore, MemorySnapshot
|
|
|
|
|
|
class TestMemoryFileBasicIO:
|
|
"""MemoryFile 基本读写测试."""
|
|
|
|
def test_read_nonexistent_file_returns_empty(self, tmp_path: Path):
|
|
mf = MemoryFile(tmp_path / "no_such.md")
|
|
assert mf.read() == ""
|
|
|
|
def test_write_and_read_back(self, tmp_path: Path):
|
|
mf = MemoryFile(tmp_path / "test.md")
|
|
mf.write("hello world")
|
|
assert mf.read() == "hello world"
|
|
|
|
def test_write_creates_parent_dirs(self, tmp_path: Path):
|
|
mf = MemoryFile(tmp_path / "deep" / "nested" / "test.md")
|
|
mf.write("content")
|
|
assert mf.read() == "content"
|
|
|
|
def test_overwrite_existing(self, tmp_path: Path):
|
|
mf = MemoryFile(tmp_path / "test.md")
|
|
mf.write("first")
|
|
mf.write("second")
|
|
assert mf.read() == "second"
|
|
|
|
|
|
class TestMemoryFileSections:
|
|
"""MemoryFile section 级别操作测试."""
|
|
|
|
def _make_file(self, tmp_path: Path, content: str) -> MemoryFile:
|
|
mf = MemoryFile(tmp_path / "test.md")
|
|
mf.write(content)
|
|
return mf
|
|
|
|
def test_read_section_from_empty_file(self, tmp_path: Path):
|
|
mf = MemoryFile(tmp_path / "empty.md")
|
|
assert mf.read_section("身份") == ""
|
|
|
|
def test_read_section_returns_content(self, tmp_path: Path):
|
|
mf = self._make_file(tmp_path, "## 身份\n我是小王\n## 性格\n友好耐心")
|
|
assert mf.read_section("身份") == "我是小王"
|
|
|
|
def test_read_section_not_found_returns_empty(self, tmp_path: Path):
|
|
mf = self._make_file(tmp_path, "## 身份\n我是小王")
|
|
assert mf.read_section("不存在") == ""
|
|
|
|
def test_add_section_creates_new(self, tmp_path: Path):
|
|
mf = self._make_file(tmp_path, "## 身份\n我是小王")
|
|
mf.add_section("性格", "友好耐心")
|
|
assert mf.read_section("性格") == "友好耐心"
|
|
assert mf.read_section("身份") == "我是小王"
|
|
|
|
def test_add_section_appends_to_existing(self, tmp_path: Path):
|
|
mf = self._make_file(tmp_path, "## 身份\n我是小王")
|
|
mf.add_section("身份", "也是AI助手")
|
|
content = mf.read_section("身份")
|
|
assert "我是小王" in content
|
|
assert "也是AI助手" in content
|
|
|
|
def test_replace_section_text(self, tmp_path: Path):
|
|
mf = self._make_file(tmp_path, "## 身份\n我是小王\n## 性格\n友好耐心")
|
|
mf.replace_section("身份", "我是小王", "我是大王")
|
|
assert mf.read_section("身份") == "我是大王"
|
|
assert mf.read_section("性格") == "友好耐心"
|
|
|
|
def test_replace_section_old_not_found_returns_false(self, tmp_path: Path):
|
|
mf = self._make_file(tmp_path, "## 身份\n我是小王")
|
|
result = mf.replace_section("身份", "不存在", "新内容")
|
|
assert result is False
|
|
|
|
def test_remove_section(self, tmp_path: Path):
|
|
mf = self._make_file(tmp_path, "## 身份\n我是小王\n## 性格\n友好耐心")
|
|
mf.remove_section("身份")
|
|
assert mf.read_section("身份") == ""
|
|
assert mf.read_section("性格") == "友好耐心"
|
|
|
|
def test_remove_nonexistent_section_no_error(self, tmp_path: Path):
|
|
mf = self._make_file(tmp_path, "## 身份\n我是小王")
|
|
mf.remove_section("不存在") # 不抛异常
|
|
assert mf.read_section("身份") == "我是小王"
|
|
|
|
def test_list_sections(self, tmp_path: Path):
|
|
mf = self._make_file(tmp_path, "## 身份\n我是小王\n## 性格\n友好耐心")
|
|
sections = mf.list_sections()
|
|
assert sections == ["身份", "性格"]
|
|
|
|
|
|
class TestMemoryFileCapacity:
|
|
"""MemoryFile 容量管理测试."""
|
|
|
|
def test_trim_to_budget_keeps_content_within_limit(self, tmp_path: Path):
|
|
mf = MemoryFile(tmp_path / "test.md", char_budget=20)
|
|
mf.write("## 身份\n我是小王一个专业的AI助手") # 超过 20 字符
|
|
mf.trim_to_budget()
|
|
content = mf.read()
|
|
assert len(content) <= 20
|
|
|
|
def test_trim_preserves_earlier_sections(self, tmp_path: Path):
|
|
mf = MemoryFile(tmp_path / "test.md", char_budget=30)
|
|
mf.write("## 身份\n我是小王\n## 性格\n友好耐心注重细节") # 性格部分超限
|
|
mf.trim_to_budget()
|
|
content = mf.read()
|
|
assert "身份" in content # 保留前面的 section
|
|
|
|
def test_no_trim_when_within_budget(self, tmp_path: Path):
|
|
mf = MemoryFile(tmp_path / "test.md", char_budget=1000)
|
|
mf.write("## 身份\n我是小王")
|
|
mf.trim_to_budget()
|
|
assert mf.read() == "## 身份\n我是小王"
|
|
|
|
def test_write_auto_trims(self, tmp_path: Path):
|
|
mf = MemoryFile(tmp_path / "test.md", char_budget=15)
|
|
mf.write("## 身份\n我是小王一个专业的AI助手非常长")
|
|
content = mf.read()
|
|
assert len(content) <= 15
|
|
|
|
|
|
class TestMemoryStoreInit:
|
|
"""MemoryStore 初始化测试."""
|
|
|
|
def test_init_creates_base_dir(self, tmp_path: Path):
|
|
store = MemoryStore(base_dir=tmp_path / "new_dir")
|
|
assert (tmp_path / "new_dir").exists()
|
|
|
|
def test_init_creates_memories_subdir(self, tmp_path: Path):
|
|
store = MemoryStore(base_dir=tmp_path)
|
|
assert (tmp_path / "memories").exists()
|
|
|
|
def test_init_creates_daily_subdir(self, tmp_path: Path):
|
|
store = MemoryStore(base_dir=tmp_path)
|
|
assert (tmp_path / "memories" / "daily").exists()
|
|
|
|
|
|
class TestMemoryStoreLoadAll:
|
|
"""MemoryStore load_all 测试."""
|
|
|
|
def test_load_all_returns_snapshot(self, tmp_path: Path):
|
|
store = MemoryStore(base_dir=tmp_path)
|
|
snapshot = store.load_all()
|
|
assert isinstance(snapshot, MemorySnapshot)
|
|
|
|
def test_load_all_empty_when_no_files(self, tmp_path: Path):
|
|
store = MemoryStore(base_dir=tmp_path)
|
|
snapshot = store.load_all()
|
|
assert snapshot.is_empty()
|
|
|
|
def test_load_all_with_content(self, tmp_path: Path):
|
|
store = MemoryStore(base_dir=tmp_path)
|
|
store.get_file("soul").write("## 身份\n我是小王")
|
|
store.get_file("user").write("## 称呼\n叫我老板")
|
|
snapshot = store.load_all()
|
|
assert "小王" in snapshot.soul
|
|
assert "老板" in snapshot.user
|
|
assert snapshot.total_chars > 0
|
|
|
|
|
|
class TestMemoryStoreBuildPrompt:
|
|
"""MemoryStore build_system_prompt 测试."""
|
|
|
|
def test_build_prompt_injects_all_sections(self, tmp_path: Path):
|
|
store = MemoryStore(base_dir=tmp_path)
|
|
store.get_file("soul").write("## 身份\n我是小王")
|
|
store.get_file("user").write("## 称呼\n叫我老板")
|
|
snapshot = store.load_all()
|
|
prompt = store.build_system_prompt(snapshot, "Be helpful.")
|
|
assert "<agent-identity>" in prompt
|
|
assert "小王" in prompt
|
|
assert "<user-profile>" in prompt
|
|
assert "老板" in prompt
|
|
assert "Be helpful." in prompt
|
|
|
|
def test_build_prompt_no_memory_returns_base_only(self, tmp_path: Path):
|
|
store = MemoryStore(base_dir=tmp_path)
|
|
snapshot = store.load_all()
|
|
prompt = store.build_system_prompt(snapshot, "Be helpful.")
|
|
assert prompt == "Be helpful."
|
|
|
|
def test_build_prompt_empty_base_with_memory(self, tmp_path: Path):
|
|
store = MemoryStore(base_dir=tmp_path)
|
|
store.get_file("soul").write("## 身份\n我是小王")
|
|
snapshot = store.load_all()
|
|
prompt = store.build_system_prompt(snapshot)
|
|
assert "<agent-identity>" in prompt
|
|
assert "小王" in prompt
|
|
|
|
|
|
class TestMemoryStoreDefaults:
|
|
"""MemoryStore ensure_defaults 测试."""
|
|
|
|
def test_ensure_defaults_creates_soul(self, tmp_path: Path):
|
|
store = MemoryStore(base_dir=tmp_path)
|
|
store.ensure_defaults()
|
|
soul = store.get_file("soul").read()
|
|
assert "AgentKit" in soul
|
|
|
|
def test_ensure_defaults_no_overwrite(self, tmp_path: Path):
|
|
store = MemoryStore(base_dir=tmp_path)
|
|
store.get_file("soul").write("## 身份\n自定义内容")
|
|
store.ensure_defaults()
|
|
soul = store.get_file("soul").read()
|
|
assert "自定义内容" in soul
|
|
assert "AgentKit" not in soul
|
|
|
|
|
|
class TestMemoryStoreDailyLogs:
|
|
"""MemoryStore 日志管理测试."""
|
|
|
|
def test_load_daily_logs_empty(self, tmp_path: Path):
|
|
store = MemoryStore(base_dir=tmp_path)
|
|
assert store.load_daily_logs() == ""
|
|
|
|
def test_load_daily_logs_with_today(self, tmp_path: Path):
|
|
store = MemoryStore(base_dir=tmp_path)
|
|
today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
|
daily_file = MemoryFile(tmp_path / "memories" / "daily" / f"{today}.md")
|
|
daily_file.write("讨论了项目架构")
|
|
logs = store.load_daily_logs()
|
|
assert "讨论了项目架构" in logs
|
|
|
|
def test_archive_old_dailies(self, tmp_path: Path):
|
|
store = MemoryStore(base_dir=tmp_path)
|
|
# 创建一个旧日志
|
|
old_date = (datetime.now(timezone.utc) - timedelta(days=5)).strftime("%Y-%m-%d")
|
|
old_file = tmp_path / "memories" / "daily" / f"{old_date}.md"
|
|
old_file.parent.mkdir(parents=True, exist_ok=True)
|
|
old_file.write_text("旧日志", encoding="utf-8")
|
|
count = store.archive_old_dailies(keep_days=2)
|
|
assert count == 1
|
|
assert not old_file.exists()
|
|
|
|
def test_get_file_daily_returns_today(self, tmp_path: Path):
|
|
store = MemoryStore(base_dir=tmp_path)
|
|
daily = store.get_file("daily")
|
|
today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
|
assert today in str(daily.path)
|
|
|
|
def test_get_file_invalid_key_raises(self, tmp_path: Path):
|
|
store = MemoryStore(base_dir=tmp_path)
|
|
with pytest.raises(ValueError, match="Invalid file_key"):
|
|
store.get_file("invalid")
|