deer-flow/backend/tests/test_memory_storage.py

"""Tests for memory storage providers (DI: FileMemoryStorage(config) / create_storage)."""

import json
import threading
from unittest.mock import patch

import pytest

from deerflow.agents.memory.backends.deermem.deermem.config import DeerMemConfig
from deerflow.agents.memory.backends.deermem.deermem.core import markdown_format as mf
from deerflow.agents.memory.backends.deermem.deermem.core.markdown_storage import MarkdownMemoryStorage
from deerflow.agents.memory.backends.deermem.deermem.core.paths import validate_agent_name
from deerflow.agents.memory.backends.deermem.deermem.core.storage import (
    FileMemoryStorage,
    MemoryStorage,
    create_empty_memory,
    create_storage,
)


def _storage_at(memory_file) -> FileMemoryStorage:
    """A FileMemoryStorage whose absolute storage_path is a single shared file."""
    resolved = str(memory_file.resolve())
    return FileMemoryStorage(DeerMemConfig(storage_path=resolved))


class TestCreateEmptyMemory:
    """Test create_empty_memory function."""

    def test_returns_valid_structure(self):
        memory = create_empty_memory()
        assert isinstance(memory, dict)
        assert memory["version"] == "1.0"
        assert "lastUpdated" in memory
        assert isinstance(memory["user"], dict)
        assert isinstance(memory["history"], dict)
        assert isinstance(memory["facts"], list)


class TestMemoryStorageInterface:
    """Test MemoryStorage abstract base class."""

    def test_abstract_methods(self):
        class TestStorage(MemoryStorage):
            pass

        with pytest.raises(TypeError):
            TestStorage(DeerMemConfig())


class TestFileMemoryStorage:
    """Test FileMemoryStorage implementation (DI: constructed with a config)."""

    def test_get_memory_file_path_global(self, tmp_path, monkeypatch):
        """DEERMEM_DATA_DIR as root + empty storage_path => global legacy path."""
        monkeypatch.setenv("DEERMEM_DATA_DIR", str(tmp_path))
        storage = FileMemoryStorage(DeerMemConfig())
        assert storage._get_memory_file_path(None) == tmp_path / "memory.json"

    def test_get_memory_file_path_agent(self, tmp_path, monkeypatch):
        """Agent facts share the user's global summary JSON path."""
        monkeypatch.setenv("DEERMEM_DATA_DIR", str(tmp_path))
        storage = FileMemoryStorage(DeerMemConfig())
        path = storage._get_memory_file_path("test-agent")
        assert path == tmp_path / "memory.json"

    @pytest.mark.parametrize("invalid_name", ["", "../etc/passwd", "agent/name", "agent\\name", "agent name", "agent@123", "agent_name"])
    def test_validate_agent_name_invalid(self, invalid_name):
        """Should raise ValueError for invalid agent names."""
        with pytest.raises(ValueError, match="Invalid agent name|Agent name must be a non-empty string"):
            validate_agent_name(invalid_name)

    def test_load_creates_empty_memory(self, tmp_path, monkeypatch):
        monkeypatch.setenv("DEERMEM_DATA_DIR", str(tmp_path))
        storage = FileMemoryStorage(DeerMemConfig())
        memory = storage.load()
        assert isinstance(memory, dict)
        assert memory["version"] == "1.0"

    def test_save_writes_to_file(self, tmp_path, monkeypatch):
        monkeypatch.setenv("DEERMEM_DATA_DIR", str(tmp_path))
        memory_file = tmp_path / "memory.json"
        storage = FileMemoryStorage(DeerMemConfig())
        result = storage.save({"version": "1.0", "facts": []})
        assert result is True
        assert memory_file.exists()
        assert "facts" not in memory_file.read_text(encoding="utf-8")

    def test_save_does_not_mutate_caller_dict(self, tmp_path, monkeypatch):
        monkeypatch.setenv("DEERMEM_DATA_DIR", str(tmp_path))
        storage = FileMemoryStorage(DeerMemConfig())
        original = {"version": "1.0", "facts": []}
        before_keys = set(original.keys())
        storage.save(original)
        assert set(original.keys()) == before_keys, "save() must not add keys to caller's dict"
        assert "lastUpdated" not in original

    def test_cache_not_corrupted_when_save_fails(self, tmp_path, monkeypatch):
        """Cache must remain clean when save() raises OSError."""
        monkeypatch.setenv("DEERMEM_DATA_DIR", str(tmp_path))
        memory_file = tmp_path / "memory.json"
        memory_file.parent.mkdir(parents=True, exist_ok=True)
        import json as _json

        memory_file.write_text(_json.dumps({"version": "2.0", "user": {"workContext": {"summary": "original"}}, "history": {}}))
        storage = FileMemoryStorage(DeerMemConfig())
        cached = storage.load()
        assert cached["user"]["workContext"]["summary"] == "original"

        with patch("builtins.open", side_effect=OSError("disk full")):
            changed = create_empty_memory()
            changed["user"]["workContext"]["summary"] = "mutated"
            result = storage.save(changed)
        assert result is False
        after = storage.load()
        assert after["user"]["workContext"]["summary"] == "original"

    def test_cache_thread_safety(self, tmp_path, monkeypatch):
        monkeypatch.setenv("DEERMEM_DATA_DIR", str(tmp_path))
        memory_file = tmp_path / "memory.json"
        memory_file.parent.mkdir(parents=True, exist_ok=True)
        import json as _json

        memory_file.write_text(_json.dumps({"version": "1.0", "facts": []}))
        storage = FileMemoryStorage(DeerMemConfig())
        errors: list[Exception] = []

        def load_many(s: FileMemoryStorage) -> None:
            try:
                for _ in range(50):
                    s.load()
            except Exception as exc:
                errors.append(exc)

        threads = [threading.Thread(target=load_many, args=(storage,)) for _ in range(8)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()
        assert not errors, f"Thread-safety errors: {errors}"

    def test_reload_forces_cache_invalidation(self, tmp_path, monkeypatch):
        monkeypatch.setenv("DEERMEM_DATA_DIR", str(tmp_path))
        memory_file = tmp_path / "memory.json"
        memory_file.parent.mkdir(parents=True, exist_ok=True)
        memory_file.write_text('{"version": "2.0", "user": {"workContext": {"summary": "initial"}}, "history": {}}')
        storage = FileMemoryStorage(DeerMemConfig())
        memory1 = storage.load()
        assert memory1["user"]["workContext"]["summary"] == "initial"
        memory_file.write_text('{"version": "2.0", "user": {"workContext": {"summary": "updated"}}, "history": {}}')
        memory2 = storage.reload()
        assert memory2["user"]["workContext"]["summary"] == "updated"


class TestCreateStorage:
    """Test create_storage(config) (replaces the old get_memory_storage() singleton)."""

    def test_returns_file_memory_storage_by_default(self):
        """Empty storage_class (default) -> FileMemoryStorage directly."""
        assert isinstance(create_storage(DeerMemConfig()), FileMemoryStorage)

    def test_raises_on_unresolvable_storage_class(self):
        """An unimportable storage_class raises ValueError (fail-fast), not a silent
        FileMemoryStorage fallback -- memory is persistent state, so a wrong store
        is a data-integrity footgun. Mirrors the manager_class resolution policy."""
        with pytest.raises(ValueError, match="storage_class"):
            create_storage(DeerMemConfig(storage_class="non.existent.StorageClass"))

    def test_raises_on_non_class_storage_class(self):
        """A storage_class that resolves to a non-class (e.g. a function) raises
        ValueError, not a silent fallback to FileMemoryStorage."""
        with pytest.raises(ValueError, match="storage_class"):
            create_storage(DeerMemConfig(storage_class="os.path.join"))

    def test_raises_on_non_subclass_storage_class(self):
        """A storage_class that is not a MemoryStorage subclass raises ValueError,
        not a silent fallback to FileMemoryStorage."""
        with pytest.raises(ValueError, match="storage_class"):
            create_storage(DeerMemConfig(storage_class="builtins.dict"))

    def test_dotted_storage_class_resolves(self):
        storage = create_storage(DeerMemConfig(storage_class="deerflow.agents.memory.backends.deermem.deermem.core.storage.FileMemoryStorage"))
        assert isinstance(storage, FileMemoryStorage)


class TestMarkdownMemoryStorage:
    """Opt-in ``storage_class="markdown"``: tolerant load path (issue #3124)."""

    def _markdown_storage_at(self, memory_file) -> MarkdownMemoryStorage:
        return MarkdownMemoryStorage(DeerMemConfig(storage_path=str(memory_file.parent.resolve())))

    def test_markdown_alias_resolves(self):
        storage = create_storage(DeerMemConfig(storage_class="markdown"))
        assert isinstance(storage, MarkdownMemoryStorage)
        assert isinstance(storage, FileMemoryStorage)

    @pytest.mark.parametrize(("version", "expected_revision"), [("1.0", 8), ("2.0", 7)])
    def test_corrupt_json_with_fenced_block_recovers(self, tmp_path, version, expected_revision):
        """A partially written JSON summary that still carries a fenced
        ```memory-json block is recovered losslessly instead of crashing."""
        memory_file = tmp_path / "memory.json"
        manifest = create_empty_memory()
        manifest["version"] = version
        if version == "2.0":
            manifest.pop("facts")  # v2 manifests keep facts in separate files.
        manifest["revision"] = 7
        manifest["user"]["workContext"] = {"summary": "recovered"}
        fenced = '{"version": 1, "revision": 7, "user": {"lang": "zh"}\n```memory-json\n' + json.dumps(manifest, ensure_ascii=False) + "\n```"
        memory_file.write_text(fenced, encoding="utf-8")
        storage = self._markdown_storage_at(memory_file)
        loaded = storage.load()
        # Legacy manifests migrate to v2 on load and advance the revision.
        assert loaded["version"] == "2.0"
        assert loaded["revision"] == expected_revision
        assert loaded["user"]["workContext"]["summary"] == "recovered"

    def test_hand_edited_markdown_without_fence_is_quarantined(self, tmp_path):
        """Markdown without a fenced block cannot be mapped onto the manifest
        schema losslessly: load() must not crash AND must not return an
        invalid shape -- the file is quarantined so nothing is silently lost."""
        memory_file = tmp_path / "memory.json"
        body = "# DeerFlow Memory\n\n- version: 2\n- revision: 5\n\n## User\n- summary: likes tea\n"
        memory_file.write_text(body, encoding="utf-8")
        storage = self._markdown_storage_at(memory_file)
        loaded = storage.load()  # must not raise
        assert isinstance(loaded, dict)
        quarantined = list(tmp_path.glob("memory.json.corrupt-*"))
        assert len(quarantined) == 1, "unreadable file must be preserved via quarantine"
        assert quarantined[0].read_text(encoding="utf-8") == body

    def test_truncated_json_is_quarantined_before_rebuild(self, tmp_path):
        """A truncated manifest must not be silently erased by the next save:
        the unreadable file is quarantined (revision reset is then safe)."""
        memory_file = tmp_path / "memory.json"
        truncated = '{"version": "2.0", "revision": 41, "user": {"work'
        memory_file.write_text(truncated, encoding="utf-8")
        storage = self._markdown_storage_at(memory_file)
        loaded = storage.load()
        assert isinstance(loaded, dict)
        assert storage.save(create_empty_memory()) is True
        quarantined = list(tmp_path.glob("memory.json.corrupt-*"))
        assert len(quarantined) == 1
        assert quarantined[0].read_text(encoding="utf-8") == truncated

    def test_fenced_block_containing_backticks_parses_losslessly(self):
        """Remembered code snippets must not terminate the JSON block."""
        manifest = create_empty_memory()
        manifest["user"]["workContext"] = {"summary": "prefers ```python\nprint('hi')\n``` snippets"}
        rendered = "# DeerFlow Memory\n\n```memory-json\n" + json.dumps(manifest, ensure_ascii=False) + "\n```\n"
        parsed = mf._parse_markdown_memory(rendered)
        assert parsed == manifest

    @pytest.mark.parametrize("newline", ["\n", "\r\n"])
    def test_fenced_summary_with_trailing_code_block_loads_without_quarantine(self, tmp_path, newline):
        manifest = create_empty_memory()
        manifest["version"] = "2.0"
        manifest.pop("facts")
        manifest["revision"] = 7
        manifest["user"]["workContext"]["summary"] = "prefers ```python snippets```"
        body = newline.join(["# Memory", "```memory-json", json.dumps(manifest, indent=2), "", "```", "Notes:", "```python", "print('example')", "```", ""])
        memory_file = tmp_path / "memory.json"
        memory_file.write_text(body, encoding="utf-8")
        assert mf._parse_markdown_memory(body) == manifest
        loaded = self._markdown_storage_at(memory_file).load()
        assert loaded["revision"] == 7
        assert loaded["user"]["workContext"] == manifest["user"]["workContext"]
        assert not list(tmp_path.glob("memory.json.corrupt-*"))

    @pytest.mark.parametrize(
        "body",
        [
            "# Memory without a fence",
            '```memory-json\n{"user": ',
            '```memory-json\n{"user": {}}',
            '```memory-json\n{"user": {}} trailing garbage\n```',
            '```memory-json\n{"user": {}}\n```python\nnotes\n```',
            "```memory-json\n[]\n```",
            "```memory-json\nnull\n```",
        ],
    )
    def test_invalid_or_non_object_fenced_summary_rejected(self, body):
        assert mf._parse_markdown_memory(body) is None