mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
d142aee83f
* fix(memory): keep sample import backups unique per run * docs(memory): clarify backup retention and align script format --------- Co-authored-by: Kiko0301 <191240073+Kiko0301@users.noreply.github.com>
263 lines
9.1 KiB
Python
263 lines
9.1 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import json
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
SCRIPT_PATH = Path(__file__).resolve().parents[2] / "scripts" / "load_memory_sample.py"
|
|
SPEC = importlib.util.spec_from_file_location("load_memory_sample", SCRIPT_PATH)
|
|
assert SPEC is not None and SPEC.loader is not None
|
|
loader = importlib.util.module_from_spec(SPEC)
|
|
SPEC.loader.exec_module(loader)
|
|
|
|
|
|
def test_parse_args_requires_one_target_mode(tmp_path):
|
|
with pytest.raises(SystemExit):
|
|
loader.parse_args(tmp_path, [])
|
|
|
|
|
|
def test_parse_args_rejects_target_with_all_users(tmp_path):
|
|
with pytest.raises(SystemExit):
|
|
loader.parse_args(tmp_path, ["--target", "memory.json", "--all-users"])
|
|
|
|
|
|
def test_parse_args_accepts_all_users(tmp_path):
|
|
args = loader.parse_args(tmp_path, ["--all-users"])
|
|
|
|
assert args.all_users is True
|
|
assert args.target is None
|
|
|
|
|
|
def test_load_sample_for_users_backs_up_every_user_before_import(tmp_path):
|
|
events = []
|
|
current = {
|
|
"u1": {"facts": [{"id": "old-1"}]},
|
|
"u2": {"facts": [{"id": "old-2"}]},
|
|
}
|
|
|
|
def load_memory(*, user_id):
|
|
events.append(("load", user_id))
|
|
return current[user_id]
|
|
|
|
def import_memory(sample, *, user_id):
|
|
assert (tmp_path / f"{user_id}.json").exists()
|
|
assert all((tmp_path / f"{uid}.json").exists() for uid in current)
|
|
events.append(("import", user_id))
|
|
return sample
|
|
|
|
count = loader.load_sample_for_users(
|
|
{"facts": [{"id": "sample"}]},
|
|
["u1", "u2"],
|
|
backup_root=tmp_path,
|
|
no_backup=False,
|
|
load_memory=load_memory,
|
|
import_memory=import_memory,
|
|
)
|
|
|
|
assert count == 2
|
|
assert events == [
|
|
("load", "u1"),
|
|
("load", "u2"),
|
|
("import", "u1"),
|
|
("import", "u2"),
|
|
]
|
|
|
|
|
|
def test_load_sample_for_users_with_no_users_is_a_noop(tmp_path):
|
|
assert (
|
|
loader.load_sample_for_users(
|
|
{"facts": []},
|
|
[],
|
|
backup_root=tmp_path,
|
|
no_backup=False,
|
|
load_memory=lambda **_: pytest.fail("unexpected load"),
|
|
import_memory=lambda *_args, **_kwargs: pytest.fail("unexpected import"),
|
|
)
|
|
== 0
|
|
)
|
|
|
|
|
|
def test_load_sample_for_users_stops_after_import_failure(tmp_path):
|
|
imported = []
|
|
|
|
def fail_second(_sample, *, user_id):
|
|
imported.append(user_id)
|
|
if user_id == "u2":
|
|
raise OSError("boom")
|
|
|
|
with pytest.raises(OSError, match="boom"):
|
|
loader.load_sample_for_users(
|
|
{"facts": []},
|
|
["u1", "u2", "u3"],
|
|
backup_root=tmp_path,
|
|
no_backup=True,
|
|
load_memory=lambda **_: {},
|
|
import_memory=fail_second,
|
|
)
|
|
|
|
assert imported == ["u1", "u2"]
|
|
|
|
|
|
def test_load_sample_for_users_rejects_unpersisted_sample(tmp_path):
|
|
with pytest.raises(OSError, match="not persisted for user u1"):
|
|
loader.load_sample_for_users(
|
|
{"facts": [{"id": "sample"}]},
|
|
["u1"],
|
|
backup_root=tmp_path,
|
|
no_backup=True,
|
|
load_memory=lambda **_: {"facts": []},
|
|
import_memory=lambda *_args, **_kwargs: {"facts": []},
|
|
)
|
|
|
|
|
|
def test_target_backups_keep_both_copies_when_imports_share_a_second(monkeypatch, tmp_path):
|
|
source = tmp_path / "sample.json"
|
|
target = tmp_path / "memory.json"
|
|
source.write_text(json.dumps({"facts": [{"id": "sample"}]}), encoding="utf-8")
|
|
target.write_text(json.dumps({"facts": [{"id": "original"}]}), encoding="utf-8")
|
|
monkeypatch.setattr(loader, "datetime", SimpleNamespace(now=lambda: datetime(2026, 10, 24, 12, 0, 0)))
|
|
|
|
args = ["--source", str(source), "--target", str(target)]
|
|
assert loader.main(args) == 0
|
|
first_backup = next(tmp_path.glob("memory.json.bak-*"))
|
|
assert json.loads(first_backup.read_text(encoding="utf-8"))["facts"][0]["id"] == "original"
|
|
|
|
assert loader.main(args) == 0
|
|
backups = list(tmp_path.glob("memory.json.bak-*"))
|
|
assert len(backups) == 2
|
|
assert json.loads(first_backup.read_text(encoding="utf-8"))["facts"][0]["id"] == "original"
|
|
second_backup = next(path for path in backups if path != first_backup)
|
|
assert json.loads(second_backup.read_text(encoding="utf-8"))["facts"][0]["id"] == "sample"
|
|
|
|
|
|
def test_target_backup_removes_reserved_file_if_copy_fails(monkeypatch, tmp_path):
|
|
target = tmp_path / "memory.json"
|
|
target.write_text("original", encoding="utf-8")
|
|
|
|
def fail_copy(_source, _destination):
|
|
raise OSError("copy failed")
|
|
|
|
monkeypatch.setattr(loader.shutil, "copy2", fail_copy)
|
|
with pytest.raises(OSError, match="copy failed"):
|
|
loader.create_target_backup(target)
|
|
|
|
assert list(tmp_path.glob("memory.json.bak-*")) == []
|
|
assert target.read_text(encoding="utf-8") == "original"
|
|
|
|
|
|
def test_load_sample_for_all_users_uses_runtime_config_resolution(monkeypatch, tmp_path):
|
|
import asyncio
|
|
|
|
import deerflow.config.app_config as app_config
|
|
|
|
config_arguments = []
|
|
|
|
def from_file(config_path=None):
|
|
config_arguments.append(config_path)
|
|
return SimpleNamespace(database=SimpleNamespace(backend="memory"))
|
|
|
|
monkeypatch.setenv("DEER_FLOW_CONFIG_PATH", str(tmp_path / "review-config.yaml"))
|
|
monkeypatch.setattr(app_config.AppConfig, "from_file", from_file)
|
|
|
|
with pytest.raises(SystemExit, match="sqlite or postgres"):
|
|
asyncio.run(
|
|
loader.load_sample_for_all_users(SCRIPT_PATH.parents[1], {"facts": []}, no_backup=True),
|
|
)
|
|
|
|
assert config_arguments == [None]
|
|
|
|
|
|
def test_load_sample_for_all_users_uses_configured_memory_manager(monkeypatch, tmp_path):
|
|
import asyncio
|
|
|
|
import app.gateway.auth.repositories.sqlite as sqlite_repository
|
|
import deerflow.agents.memory.manager as memory_manager
|
|
import deerflow.config.app_config as app_config
|
|
import deerflow.config.paths as config_paths
|
|
import deerflow.persistence.engine as persistence_engine
|
|
|
|
manager_factory_calls = 0
|
|
loaded = []
|
|
imported = []
|
|
current = {user_id: {"facts": [{"id": f"old-{user_id}"}]} for user_id in ("u1", "u2")}
|
|
|
|
class FakeMemoryManager:
|
|
def get_memory(self, *, user_id):
|
|
loaded.append(user_id)
|
|
return current[user_id]
|
|
|
|
def import_memory(self, sample, *, user_id):
|
|
imported.append((user_id, sample))
|
|
current[user_id] = sample
|
|
return sample
|
|
|
|
manager = FakeMemoryManager()
|
|
|
|
def get_memory_manager():
|
|
nonlocal manager_factory_calls
|
|
manager_factory_calls += 1
|
|
return manager
|
|
|
|
class FakeUserRepository:
|
|
def __init__(self, _session_factory):
|
|
pass
|
|
|
|
async def list_user_ids(self):
|
|
return ["u1", "u2"]
|
|
|
|
async def init_engine_from_config(_database):
|
|
pass
|
|
|
|
async def close_engine():
|
|
pass
|
|
|
|
monkeypatch.setattr(memory_manager, "get_memory_manager", get_memory_manager)
|
|
monkeypatch.setattr(sqlite_repository, "SQLiteUserRepository", FakeUserRepository)
|
|
monkeypatch.setattr(
|
|
app_config.AppConfig,
|
|
"from_file",
|
|
lambda *_args: SimpleNamespace(database=SimpleNamespace(backend="sqlite")),
|
|
)
|
|
monkeypatch.setattr(config_paths, "get_paths", lambda: SimpleNamespace(base_dir=tmp_path))
|
|
monkeypatch.setattr(persistence_engine, "init_engine_from_config", init_engine_from_config)
|
|
monkeypatch.setattr(persistence_engine, "get_session_factory", object)
|
|
monkeypatch.setattr(persistence_engine, "close_engine", close_engine)
|
|
monkeypatch.setattr(loader, "datetime", SimpleNamespace(now=lambda: datetime(2026, 10, 24, 12, 0, 0)))
|
|
|
|
sample = {"facts": [{"id": "sample"}]}
|
|
count, backup_root = asyncio.run(
|
|
loader.load_sample_for_all_users(SCRIPT_PATH.parents[1], sample, no_backup=False),
|
|
)
|
|
|
|
assert count == 2
|
|
assert manager_factory_calls == 1
|
|
assert loaded == ["u1", "u2"]
|
|
assert imported == [("u1", sample), ("u2", sample)]
|
|
assert backup_root is not None
|
|
assert json.loads((backup_root / "u1.json").read_text())["facts"][0]["id"] == "old-u1"
|
|
assert json.loads((backup_root / "u2.json").read_text())["facts"][0]["id"] == "old-u2"
|
|
|
|
second_count, second_backup_root = asyncio.run(
|
|
loader.load_sample_for_all_users(SCRIPT_PATH.parents[1], sample, no_backup=False),
|
|
)
|
|
assert second_count == 2
|
|
assert second_backup_root is not None and second_backup_root != backup_root
|
|
assert json.loads((backup_root / "u1.json").read_text())["facts"][0]["id"] == "old-u1"
|
|
assert json.loads((backup_root / "u2.json").read_text())["facts"][0]["id"] == "old-u2"
|
|
assert json.loads((second_backup_root / "u1.json").read_text())["facts"][0]["id"] == "sample"
|
|
assert json.loads((second_backup_root / "u2.json").read_text())["facts"][0]["id"] == "sample"
|
|
|
|
|
|
def test_require_persistent_database_rejects_memory():
|
|
with pytest.raises(SystemExit, match="sqlite or postgres"):
|
|
loader.require_persistent_database("memory")
|
|
|
|
|
|
@pytest.mark.parametrize("backend", ["sqlite", "postgres"])
|
|
def test_require_persistent_database_accepts_persistent_backends(backend):
|
|
loader.require_persistent_database(backend)
|