refactor: add close() to Memory and thread-safe telemetry singleton
This commit is contained in:
kartik-mem0
@@ -0,0 +1,164 @@
|
||||
"""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()
|
||||
+81
-17
@@ -30,9 +30,10 @@ class TestTelemetryDisabled:
|
||||
with patch.object(telemetry_module, "MEM0_TELEMETRY", False):
|
||||
mock_instance = MagicMock()
|
||||
mock_client_telemetry = MagicMock()
|
||||
with patch.object(telemetry_module, "client_telemetry", mock_client_telemetry):
|
||||
telemetry_module.capture_client_event("test.event", mock_instance)
|
||||
mock_client_telemetry.capture_event.assert_not_called()
|
||||
telemetry_module._client_telemetry = mock_client_telemetry
|
||||
telemetry_module.capture_client_event("test.event", mock_instance)
|
||||
mock_client_telemetry.capture_event.assert_not_called()
|
||||
telemetry_module._client_telemetry = None
|
||||
|
||||
def test_instance_capture_event_noop_when_posthog_is_none(self):
|
||||
"""AnonymousTelemetry.capture_event() should be a no-op when posthog is None."""
|
||||
@@ -69,26 +70,89 @@ class TestTelemetryEnabled:
|
||||
assert at.user_id == "test-user"
|
||||
|
||||
def test_capture_event_sends_when_enabled(self):
|
||||
"""capture_event() should create AnonymousTelemetry and call capture when enabled."""
|
||||
"""capture_event() should use singleton AnonymousTelemetry and call capture when enabled."""
|
||||
with patch.object(telemetry_module, "MEM0_TELEMETRY", True):
|
||||
with patch("mem0.memory.telemetry.AnonymousTelemetry") as mock_cls:
|
||||
mock_at = MagicMock()
|
||||
mock_cls.return_value = mock_at
|
||||
mock_memory = MagicMock()
|
||||
mock_memory.config.graph_store.config = None
|
||||
mock_memory.api_version = "v1"
|
||||
telemetry_module.capture_event("test.event", mock_memory)
|
||||
mock_at.capture_event.assert_called_once()
|
||||
mock_at = MagicMock()
|
||||
telemetry_module._oss_telemetry = mock_at
|
||||
mock_memory = MagicMock()
|
||||
mock_memory.config.graph_store.config = None
|
||||
mock_memory.api_version = "v1"
|
||||
telemetry_module.capture_event("test.event", mock_memory)
|
||||
mock_at.capture_event.assert_called_once()
|
||||
telemetry_module._oss_telemetry = None
|
||||
|
||||
def test_capture_client_event_sends_when_enabled(self):
|
||||
"""capture_client_event() should call client_telemetry.capture_event when enabled."""
|
||||
with patch.object(telemetry_module, "MEM0_TELEMETRY", True):
|
||||
mock_client_telemetry = MagicMock()
|
||||
with patch.object(telemetry_module, "client_telemetry", mock_client_telemetry):
|
||||
mock_instance = MagicMock()
|
||||
mock_instance.user_email = "test@example.com"
|
||||
telemetry_module.capture_client_event("test.event", mock_instance)
|
||||
mock_client_telemetry.capture_event.assert_called_once()
|
||||
telemetry_module._client_telemetry = mock_client_telemetry
|
||||
mock_instance = MagicMock()
|
||||
mock_instance.user_email = "test@example.com"
|
||||
telemetry_module.capture_client_event("test.event", mock_instance)
|
||||
mock_client_telemetry.capture_event.assert_called_once()
|
||||
telemetry_module._client_telemetry = None
|
||||
|
||||
def test_capture_event_reuses_singleton(self):
|
||||
"""capture_event() must reuse a single AnonymousTelemetry, not create one per call."""
|
||||
with patch.object(telemetry_module, "MEM0_TELEMETRY", True):
|
||||
with patch("mem0.memory.telemetry.AnonymousTelemetry") as mock_cls:
|
||||
mock_at = MagicMock()
|
||||
mock_cls.return_value = mock_at
|
||||
telemetry_module._oss_telemetry = None
|
||||
mock_memory = MagicMock()
|
||||
mock_memory.config.graph_store.config = None
|
||||
mock_memory.api_version = "v1"
|
||||
|
||||
telemetry_module.capture_event("event1", mock_memory)
|
||||
telemetry_module.capture_event("event2", mock_memory)
|
||||
|
||||
assert mock_cls.call_count == 1
|
||||
assert mock_at.capture_event.call_count == 2
|
||||
|
||||
telemetry_module._oss_telemetry = None
|
||||
|
||||
|
||||
class TestTelemetryShutdown:
|
||||
"""Verify shutdown_telemetry() cleans up all singleton instances."""
|
||||
|
||||
def test_shutdown_calls_close_on_singletons(self):
|
||||
mock_oss = MagicMock()
|
||||
mock_client = MagicMock()
|
||||
telemetry_module._oss_telemetry = mock_oss
|
||||
telemetry_module._client_telemetry = mock_client
|
||||
|
||||
telemetry_module.shutdown_telemetry()
|
||||
|
||||
mock_oss.close.assert_called_once()
|
||||
mock_client.close.assert_called_once()
|
||||
assert telemetry_module._oss_telemetry is None
|
||||
assert telemetry_module._client_telemetry is None
|
||||
|
||||
def test_shutdown_is_safe_when_no_singletons(self):
|
||||
telemetry_module._oss_telemetry = None
|
||||
telemetry_module._client_telemetry = None
|
||||
telemetry_module.shutdown_telemetry() # should not raise
|
||||
|
||||
|
||||
class TestNoThreadLeak:
|
||||
"""Verify capture_event does not leak threads."""
|
||||
|
||||
def test_multiple_capture_events_single_posthog_instance(self):
|
||||
with patch.object(telemetry_module, "MEM0_TELEMETRY", True):
|
||||
with patch("mem0.memory.telemetry.Posthog") as mock_posthog:
|
||||
with patch("mem0.memory.telemetry.get_or_create_user_id", return_value="test"):
|
||||
telemetry_module._oss_telemetry = None
|
||||
|
||||
mock_memory = MagicMock()
|
||||
mock_memory.config.graph_store.config = None
|
||||
mock_memory.api_version = "v1"
|
||||
|
||||
for _ in range(10):
|
||||
telemetry_module.capture_event("test.event", mock_memory)
|
||||
|
||||
assert mock_posthog.call_count == 1
|
||||
|
||||
telemetry_module._oss_telemetry = None
|
||||
|
||||
|
||||
class TestTelemetryEnvVar:
|
||||
|
||||
Reference in New Issue
Block a user