mem0/tests/test_memory_lifecycle.py

"""Tests for Memory resource cleanup: close() and context manager."""
from unittest.mock import MagicMock, patch

import pytest


@pytest.fixture
def memory_instance():
    """Create a Memory instance with all dependencies mocked."""
    with (
        patch("mem0.memory.main.EmbedderFactory") as mock_embedder,
        patch("mem0.memory.main.VectorStoreFactory") as mock_vector_store,
        patch("mem0.memory.main.LlmFactory") as mock_llm,
        patch("mem0.memory.telemetry.capture_event"),
        patch("mem0.memory.main.GraphStoreFactory") as mock_graph_store,
    ):
        mock_embedder.create.return_value = MagicMock()
        mock_embedder.create.return_value.config.embedding_dims = 1536
        mock_vector_store.create.return_value = MagicMock()
        mock_llm.create.return_value = MagicMock()
        mock_graph_store.create.return_value = MagicMock()

        from mem0.memory.main import Memory
        m = Memory()
        yield m


class TestMemoryClose:
    """Verify Memory.close() releases resources."""

    def test_close_shuts_down_telemetry(self, memory_instance):
        with patch("mem0.memory.telemetry.shutdown_telemetry") as mock_shutdown:
            memory_instance.close()
            mock_shutdown.assert_called_once()

    def test_close_closes_db(self, memory_instance):
        memory_instance.db = MagicMock()
        memory_instance.close()
        memory_instance.db.close.assert_called_once()

    def test_close_closes_vector_store_if_available(self, memory_instance):
        mock_vs = MagicMock(spec=["close"])
        memory_instance.vector_store = mock_vs
        memory_instance.close()
        mock_vs.close.assert_called_once()

    def test_close_skips_vector_store_without_close(self, memory_instance):
        mock_vs = MagicMock(spec=[])
        memory_instance.vector_store = mock_vs
        memory_instance.close()  # should not raise

    def test_close_is_idempotent(self, memory_instance):
        memory_instance.close()
        memory_instance.close()  # should not raise

    def test_close_continues_on_db_error(self, memory_instance):
        memory_instance.db = MagicMock()
        memory_instance.db.close.side_effect = Exception("db error")
        memory_instance.close()  # should not raise


class TestMemoryContextManager:
    """Verify Memory works as a context manager."""

    def test_enter_returns_self(self, memory_instance):
        assert memory_instance.__enter__() is memory_instance

    def test_exit_calls_close(self, memory_instance):
        memory_instance.db = MagicMock()
        with patch("mem0.memory.telemetry.shutdown_telemetry"):
            memory_instance.__exit__(None, None, None)
        memory_instance.db.close.assert_called_once()

    def test_does_not_suppress_exceptions(self, memory_instance):
        assert memory_instance.__exit__(ValueError, ValueError("x"), None) is False

    def test_with_statement(self, memory_instance):
        memory_instance.db = MagicMock()
        with patch("mem0.memory.telemetry.shutdown_telemetry"):
            with memory_instance:
                pass
        memory_instance.db.close.assert_called_once()

    def test_with_statement_on_exception(self, memory_instance):
        memory_instance.db = MagicMock()
        with patch("mem0.memory.telemetry.shutdown_telemetry"):
            with pytest.raises(RuntimeError):
                with memory_instance:
                    raise RuntimeError("boom")
        memory_instance.db.close.assert_called_once()




@pytest.fixture
def async_memory_instance():
    """Create an AsyncMemory instance with all dependencies mocked."""
    with (
        patch("mem0.memory.main.EmbedderFactory") as mock_embedder,
        patch("mem0.memory.main.VectorStoreFactory") as mock_vector_store,
        patch("mem0.memory.main.LlmFactory") as mock_llm,
        patch("mem0.memory.telemetry.capture_event"),
        patch("mem0.memory.main.GraphStoreFactory") as mock_graph_store,
    ):
        mock_embedder.create.return_value = MagicMock()
        mock_embedder.create.return_value.config.embedding_dims = 1536
        mock_vector_store.create.return_value = MagicMock()
        mock_llm.create.return_value = MagicMock()
        mock_graph_store.create.return_value = MagicMock()

        from mem0.memory.main import AsyncMemory
        m = AsyncMemory()
        yield m


class TestAsyncMemoryClose:
    """Verify AsyncMemory.close() releases resources."""

    def test_close_shuts_down_telemetry(self, async_memory_instance):
        with patch("mem0.memory.telemetry.shutdown_telemetry") as mock_shutdown:
            async_memory_instance.close()
            mock_shutdown.assert_called_once()

    def test_close_closes_db(self, async_memory_instance):
        async_memory_instance.db = MagicMock()
        async_memory_instance.close()
        async_memory_instance.db.close.assert_called_once()

    def test_close_is_idempotent(self, async_memory_instance):
        async_memory_instance.close()
        async_memory_instance.close()  # should not raise


class TestAsyncMemoryContextManager:
    """Verify AsyncMemory works as an async context manager."""

    @pytest.mark.asyncio
    async def test_aenter_returns_self(self, async_memory_instance):
        result = await async_memory_instance.__aenter__()
        assert result is async_memory_instance

    @pytest.mark.asyncio
    async def test_aexit_calls_close(self, async_memory_instance):
        async_memory_instance.db = MagicMock()
        with patch("mem0.memory.telemetry.shutdown_telemetry"):
            await async_memory_instance.__aexit__(None, None, None)
        async_memory_instance.db.close.assert_called_once()

    @pytest.mark.asyncio
    async def test_async_with_statement(self, async_memory_instance):
        async_memory_instance.db = MagicMock()
        with patch("mem0.memory.telemetry.shutdown_telemetry"):
            async with async_memory_instance:
                pass
        async_memory_instance.db.close.assert_called_once()

    @pytest.mark.asyncio
    async def test_async_with_statement_on_exception(self, async_memory_instance):
        async_memory_instance.db = MagicMock()
        with patch("mem0.memory.telemetry.shutdown_telemetry"):
            with pytest.raises(RuntimeError):
                async with async_memory_instance:
                    raise RuntimeError("boom")
        async_memory_instance.db.close.assert_called_once()