MemOS-MCP

import uuid

from unittest.mock import MagicMock, patch

import pytest

from memos.configs.memory import TreeTextMemoryConfig
from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata
from memos.memories.textual.tree import TreeTextMemory


@pytest.fixture
def mock_config():
    config = TreeTextMemoryConfig(
        extractor_llm={
            "backend": "openai",
            "config": {
                "model_name_or_path": "gpt-4o",
                "api_key": "test_api_key",
            },
        },
        dispatcher_llm={
            "backend": "openai",
            "config": {
                "model_name_or_path": "gpt-4o",
                "api_key": "test_api_key",
            },
        },
        embedder={
            "backend": "ollama",
            "config": {
                "model_name_or_path": "default",
            },
        },
        graph_db={
            "backend": "neo4j",
            "config": {
                "uri": "bolt://localhost:7687",
                "user": "neo4j",
                "password": "test_password",
                "db_name": "test",
            },
        },
        memory_filename="memory.json",
    )
    return config


@pytest.fixture
def mock_tree_text_memory(mock_config):
    with (
        patch("memos.llms.factory.LLMFactory.from_config"),
        patch("memos.embedders.factory.EmbedderFactory.from_config"),
        patch("memos.graph_dbs.factory.GraphStoreFactory.from_config"),
        patch("memos.memories.textual.tree_text_memory.organize.manager.MemoryManager"),
    ):
        instance = TreeTextMemory(mock_config)
        yield instance


def test_add_calls_manager(mock_tree_text_memory):
    mock_tree_text_memory.memory_manager.add = MagicMock()
    mock_item = TextualMemoryItem(
        id=str(uuid.uuid4()),
        memory="Test memory",
        metadata=TreeNodeTextualMemoryMetadata(updated_at=None),
    )
    mock_tree_text_memory.add([mock_item])
    mock_tree_text_memory.memory_manager.add.assert_called_once()


def test_get_working_memory_sorted(mock_tree_text_memory):
    older = TextualMemoryItem(
        id=str(uuid.uuid4()),
        memory="Older",
        metadata=TreeNodeTextualMemoryMetadata(updated_at="2020-01-01"),
    )
    newer = TextualMemoryItem(
        id=str(uuid.uuid4()),
        memory="Newer",
        metadata=TreeNodeTextualMemoryMetadata(updated_at="2025-01-01"),
    )
    mock_tree_text_memory.graph_store.get_all_memory_items = MagicMock(
        return_value=[older.model_dump(), newer.model_dump()]
    )

    result = mock_tree_text_memory.get_working_memory()
    assert result[0].id == newer.id


def test_get_memory_found(mock_tree_text_memory):
    test_id = str(uuid.uuid4())
    fake_record = {"id": test_id, "memory": "Test", "metadata": {}}
    mock_tree_text_memory.graph_store.get_node = MagicMock(return_value=fake_record)

    memory = mock_tree_text_memory.get(test_id)
    assert memory.id == test_id


def test_get_memory_not_found(mock_tree_text_memory):
    mock_tree_text_memory.graph_store.get_node = MagicMock(return_value=None)
    with pytest.raises(ValueError):
        mock_tree_text_memory.get(str(uuid.uuid4()))


def test_delete_all(mock_tree_text_memory):
    mock_tree_text_memory.graph_store.clear = MagicMock()
    mock_tree_text_memory.delete_all()
    mock_tree_text_memory.graph_store.clear.assert_called_once()


def test_load_file_not_exists(mock_tree_text_memory, tmp_path):
    mock_tree_text_memory.config.memory_filename = "memory.json"
    mock_tree_text_memory.graph_store.import_graph = MagicMock()

    result = tmp_path / "does_not_exist"
    mock_tree_text_memory.load(str(result))
    # Should log a warning but not raise


def test_dump_and_load_success(tmp_path, mock_tree_text_memory):
    mock_tree_text_memory.graph_store.export_graph = MagicMock(
        return_value={"nodes": [{"id": "1"}]}
    )
    mock_tree_text_memory.config.memory_filename = "memory.json"
    mock_tree_text_memory.dump(str(tmp_path))

    dumped_file = tmp_path / "memory.json"
    assert dumped_file.exists()


def test_drop_creates_backup_and_cleans(mock_tree_text_memory):
    mock_tree_text_memory.dump = MagicMock()
    mock_tree_text_memory._cleanup_old_backups = MagicMock()
    mock_tree_text_memory.graph_store.drop_database = MagicMock()

    mock_tree_text_memory.drop(keep_last_n=1)
    mock_tree_text_memory.dump.assert_called_once()
    mock_tree_text_memory._cleanup_old_backups.assert_called_once()
    mock_tree_text_memory.graph_store.drop_database.assert_called_once()