From dc883b0f9e2154ca3053c1f44df0ce788d709536 Mon Sep 17 00:00:00 2001 From: Karthikeya Kollu Date: Thu, 25 Sep 2025 22:30:03 +0530 Subject: [PATCH] [test] Add comprehensive test suite for SQLiteManager (#3494) --- tests/memory/test_storage.py | 282 +++++++++++++++++++++++++++++++++++ 1 file changed, 282 insertions(+) create mode 100644 tests/memory/test_storage.py diff --git a/tests/memory/test_storage.py b/tests/memory/test_storage.py new file mode 100644 index 000000000..8c191c3be --- /dev/null +++ b/tests/memory/test_storage.py @@ -0,0 +1,282 @@ +import os +import sqlite3 +import tempfile +import uuid +from datetime import datetime + +import pytest + +from mem0.memory.storage import SQLiteManager + + +class TestSQLiteManager: + """Comprehensive test cases for SQLiteManager class.""" + + @pytest.fixture + def temp_db_path(self): + """Create temporary database file.""" + temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db") + temp_db.close() + yield temp_db.name + if os.path.exists(temp_db.name): + os.unlink(temp_db.name) + + @pytest.fixture + def sqlite_manager(self, temp_db_path): + """Create SQLiteManager instance with temporary database.""" + manager = SQLiteManager(temp_db_path) + yield manager + if manager.connection: + manager.close() + + @pytest.fixture + def memory_manager(self): + """Create in-memory SQLiteManager instance.""" + manager = SQLiteManager(":memory:") + yield manager + if manager.connection: + manager.close() + + @pytest.fixture + def sample_data(self): + """Sample test data.""" + now = datetime.now().isoformat() + return { + "memory_id": str(uuid.uuid4()), + "old_memory": "Old memory content", + "new_memory": "New memory content", + "event": "ADD", + "created_at": now, + "updated_at": now, + "actor_id": "test_actor", + "role": "user", + } + + # ========== Initialization Tests ========== + + @pytest.mark.parametrize("db_type,path", [("file", "temp_db_path"), ("memory", ":memory:")]) + def test_initialization(self, db_type, path, request): + """Test SQLiteManager initialization with different database types.""" + if db_type == "file": + db_path = request.getfixturevalue(path) + else: + db_path = path + + manager = SQLiteManager(db_path) + assert manager.connection is not None + assert manager.db_path == db_path + manager.close() + + def test_table_schema_creation(self, sqlite_manager): + """Test that history table is created with correct schema.""" + cursor = sqlite_manager.connection.cursor() + cursor.execute("PRAGMA table_info(history)") + columns = {row[1] for row in cursor.fetchall()} + + expected_columns = { + "id", + "memory_id", + "old_memory", + "new_memory", + "event", + "created_at", + "updated_at", + "is_deleted", + "actor_id", + "role", + } + assert columns == expected_columns + + # ========== Add History Tests ========== + + def test_add_history_basic(self, sqlite_manager, sample_data): + """Test basic add_history functionality.""" + sqlite_manager.add_history( + memory_id=sample_data["memory_id"], + old_memory=sample_data["old_memory"], + new_memory=sample_data["new_memory"], + event=sample_data["event"], + created_at=sample_data["created_at"], + actor_id=sample_data["actor_id"], + role=sample_data["role"], + ) + + cursor = sqlite_manager.connection.cursor() + cursor.execute("SELECT * FROM history WHERE memory_id = ?", (sample_data["memory_id"],)) + result = cursor.fetchone() + + assert result is not None + assert result[1] == sample_data["memory_id"] + assert result[2] == sample_data["old_memory"] + assert result[3] == sample_data["new_memory"] + assert result[4] == sample_data["event"] + assert result[8] == sample_data["actor_id"] + assert result[9] == sample_data["role"] + + @pytest.mark.parametrize( + "old_memory,new_memory,is_deleted", [(None, "New memory", 0), ("Old memory", None, 1), (None, None, 1)] + ) + def test_add_history_optional_params(self, sqlite_manager, sample_data, old_memory, new_memory, is_deleted): + """Test add_history with various optional parameter combinations.""" + sqlite_manager.add_history( + memory_id=sample_data["memory_id"], + old_memory=old_memory, + new_memory=new_memory, + event="UPDATE", + updated_at=sample_data["updated_at"], + is_deleted=is_deleted, + actor_id=sample_data["actor_id"], + role=sample_data["role"], + ) + + cursor = sqlite_manager.connection.cursor() + cursor.execute("SELECT * FROM history WHERE memory_id = ?", (sample_data["memory_id"],)) + result = cursor.fetchone() + + assert result[2] == old_memory + assert result[3] == new_memory + assert result[6] == sample_data["updated_at"] + assert result[7] == is_deleted + + def test_add_history_generates_unique_ids(self, sqlite_manager, sample_data): + """Test that add_history generates unique IDs for each record.""" + for i in range(3): + sqlite_manager.add_history( + memory_id=sample_data["memory_id"], + old_memory=f"Memory {i}", + new_memory=f"Updated Memory {i}", + event="ADD" if i == 0 else "UPDATE", + ) + + cursor = sqlite_manager.connection.cursor() + cursor.execute("SELECT id FROM history WHERE memory_id = ?", (sample_data["memory_id"],)) + ids = [row[0] for row in cursor.fetchall()] + + assert len(ids) == 3 + assert len(set(ids)) == 3 + + # ========== Get History Tests ========== + + def test_get_history_empty(self, sqlite_manager): + """Test get_history for non-existent memory_id.""" + result = sqlite_manager.get_history("non-existent-id") + assert result == [] + + def test_get_history_single_record(self, sqlite_manager, sample_data): + """Test get_history for single record.""" + sqlite_manager.add_history( + memory_id=sample_data["memory_id"], + old_memory=sample_data["old_memory"], + new_memory=sample_data["new_memory"], + event=sample_data["event"], + created_at=sample_data["created_at"], + actor_id=sample_data["actor_id"], + role=sample_data["role"], + ) + + result = sqlite_manager.get_history(sample_data["memory_id"]) + + assert len(result) == 1 + record = result[0] + assert record["memory_id"] == sample_data["memory_id"] + assert record["old_memory"] == sample_data["old_memory"] + assert record["new_memory"] == sample_data["new_memory"] + assert record["event"] == sample_data["event"] + assert record["created_at"] == sample_data["created_at"] + assert record["actor_id"] == sample_data["actor_id"] + assert record["role"] == sample_data["role"] + assert record["is_deleted"] is False + + def test_get_history_chronological_ordering(self, sqlite_manager, sample_data): + """Test get_history returns records in chronological order.""" + import time + + timestamps = [] + for i in range(3): + ts = datetime.now().isoformat() + timestamps.append(ts) + sqlite_manager.add_history( + memory_id=sample_data["memory_id"], + old_memory=f"Memory {i}", + new_memory=f"Memory {i+1}", + event="ADD" if i == 0 else "UPDATE", + created_at=ts, + updated_at=ts if i > 0 else None, + ) + time.sleep(0.01) + + result = sqlite_manager.get_history(sample_data["memory_id"]) + result_timestamps = [r["created_at"] for r in result] + assert result_timestamps == sorted(timestamps) + + def test_migration_preserves_data(self, temp_db_path, sample_data): + """Test that migration preserves existing data.""" + manager1 = SQLiteManager(temp_db_path) + manager1.add_history( + memory_id=sample_data["memory_id"], + old_memory=sample_data["old_memory"], + new_memory=sample_data["new_memory"], + event=sample_data["event"], + created_at=sample_data["created_at"], + ) + original_data = manager1.get_history(sample_data["memory_id"]) + manager1.close() + + manager2 = SQLiteManager(temp_db_path) + migrated_data = manager2.get_history(sample_data["memory_id"]) + manager2.close() + + assert len(migrated_data) == len(original_data) + assert migrated_data[0]["memory_id"] == original_data[0]["memory_id"] + assert migrated_data[0]["new_memory"] == original_data[0]["new_memory"] + + def test_large_batch_operations(self, sqlite_manager): + """Test performance with large batch of operations.""" + batch_size = 1000 + memory_ids = [str(uuid.uuid4()) for _ in range(batch_size)] + for i, memory_id in enumerate(memory_ids): + sqlite_manager.add_history( + memory_id=memory_id, old_memory=None, new_memory=f"Batch memory {i}", event="ADD" + ) + + cursor = sqlite_manager.connection.cursor() + cursor.execute("SELECT COUNT(*) FROM history") + count = cursor.fetchone()[0] + assert count == batch_size + + for memory_id in memory_ids[:10]: + result = sqlite_manager.get_history(memory_id) + assert len(result) == 1 + + # ========== Tests for Migration, Reset, and Close ========== + + def test_explicit_old_schema_migration(self, temp_db_path): + """Test migration path from a legacy schema to new schema.""" + # Create a legacy 'history' table missing new columns + legacy_conn = sqlite3.connect(temp_db_path) + legacy_conn.execute(""" + CREATE TABLE history ( + id TEXT PRIMARY KEY, + memory_id TEXT, + old_memory TEXT, + new_memory TEXT, + event TEXT, + created_at DATETIME + ) + """) + legacy_id = str(uuid.uuid4()) + legacy_conn.execute( + "INSERT INTO history (id, memory_id, old_memory, new_memory, event, created_at) VALUES (?, ?, ?, ?, ?, ?)", + (legacy_id, "m1", "o", "n", "ADD", datetime.now().isoformat()), + ) + legacy_conn.commit() + legacy_conn.close() + + # Trigger migration + mgr = SQLiteManager(temp_db_path) + history = mgr.get_history("m1") + assert len(history) == 1 + assert history[0]["id"] == legacy_id + assert history[0]["actor_id"] is None + assert history[0]["is_deleted"] is False + mgr.close()