1+ """Tests for database migration deduplication functionality."""
2+
3+ import pytest
4+ from pathlib import Path
5+ from unittest .mock import patch , AsyncMock , call , MagicMock
6+
7+ from basic_memory .config import BasicMemoryConfig
8+ from basic_memory import db
9+
10+
11+ @pytest .fixture
12+ def mock_alembic_config ():
13+ """Mock Alembic config to avoid actual migration runs."""
14+ with patch ("basic_memory.db.Config" ) as mock_config_class :
15+ mock_config = MagicMock ()
16+ mock_config_class .return_value = mock_config
17+ yield mock_config
18+
19+
20+ @pytest .fixture
21+ def mock_alembic_command ():
22+ """Mock Alembic command to avoid actual migration runs."""
23+ with patch ("basic_memory.db.command" ) as mock_command :
24+ yield mock_command
25+
26+
27+ @pytest .fixture
28+ def mock_search_repository ():
29+ """Mock SearchRepository to avoid database dependencies."""
30+ with patch ("basic_memory.db.SearchRepository" ) as mock_repo_class :
31+ mock_repo = AsyncMock ()
32+ mock_repo_class .return_value = mock_repo
33+ yield mock_repo
34+
35+
36+ # Use the app_config fixture from conftest.py
37+
38+
39+ @pytest .mark .asyncio
40+ async def test_migration_deduplication_single_call (
41+ app_config , mock_alembic_config , mock_alembic_command , mock_search_repository
42+ ):
43+ """Test that migrations are only run once when called multiple times."""
44+ # Reset module state
45+ db ._migrations_completed = False
46+ db ._engine = None
47+ db ._session_maker = None
48+
49+ # First call should run migrations
50+ await db .run_migrations (app_config )
51+
52+ # Verify migrations were called
53+ mock_alembic_command .upgrade .assert_called_once_with (mock_alembic_config , "head" )
54+ mock_search_repository .init_search_index .assert_called_once ()
55+
56+ # Reset mocks for second call
57+ mock_alembic_command .reset_mock ()
58+ mock_search_repository .reset_mock ()
59+
60+ # Second call should skip migrations
61+ await db .run_migrations (app_config )
62+
63+ # Verify migrations were NOT called again
64+ mock_alembic_command .upgrade .assert_not_called ()
65+ mock_search_repository .init_search_index .assert_not_called ()
66+
67+
68+ @pytest .mark .asyncio
69+ async def test_migration_force_parameter (
70+ app_config , mock_alembic_config , mock_alembic_command , mock_search_repository
71+ ):
72+ """Test that migrations can be forced to run even if already completed."""
73+ # Reset module state
74+ db ._migrations_completed = False
75+ db ._engine = None
76+ db ._session_maker = None
77+
78+ # First call should run migrations
79+ await db .run_migrations (app_config )
80+
81+ # Verify migrations were called
82+ mock_alembic_command .upgrade .assert_called_once_with (mock_alembic_config , "head" )
83+ mock_search_repository .init_search_index .assert_called_once ()
84+
85+ # Reset mocks for forced call
86+ mock_alembic_command .reset_mock ()
87+ mock_search_repository .reset_mock ()
88+
89+ # Forced call should run migrations again
90+ await db .run_migrations (app_config , force = True )
91+
92+ # Verify migrations were called again
93+ mock_alembic_command .upgrade .assert_called_once_with (mock_alembic_config , "head" )
94+ mock_search_repository .init_search_index .assert_called_once ()
95+
96+
97+ @pytest .mark .asyncio
98+ async def test_migration_state_reset_on_shutdown ():
99+ """Test that migration state is reset when database is shut down."""
100+ # Set up completed state
101+ db ._migrations_completed = True
102+ db ._engine = AsyncMock ()
103+ db ._session_maker = AsyncMock ()
104+
105+ # Shutdown should reset state
106+ await db .shutdown_db ()
107+
108+ # Verify state was reset
109+ assert db ._migrations_completed is False
110+ assert db ._engine is None
111+ assert db ._session_maker is None
112+
113+
114+ @pytest .mark .asyncio
115+ async def test_get_or_create_db_runs_migrations_automatically (
116+ app_config , mock_alembic_config , mock_alembic_command , mock_search_repository
117+ ):
118+ """Test that get_or_create_db runs migrations automatically."""
119+ # Reset module state
120+ db ._migrations_completed = False
121+ db ._engine = None
122+ db ._session_maker = None
123+
124+ # First call should create engine and run migrations
125+ engine , session_maker = await db .get_or_create_db (
126+ app_config .database_path , app_config = app_config
127+ )
128+
129+ # Verify we got valid objects
130+ assert engine is not None
131+ assert session_maker is not None
132+
133+ # Verify migrations were called
134+ mock_alembic_command .upgrade .assert_called_once_with (mock_alembic_config , "head" )
135+ mock_search_repository .init_search_index .assert_called_once ()
136+
137+
138+ @pytest .mark .asyncio
139+ async def test_get_or_create_db_skips_migrations_when_disabled (
140+ app_config , mock_alembic_config , mock_alembic_command , mock_search_repository
141+ ):
142+ """Test that get_or_create_db can skip migrations when ensure_migrations=False."""
143+ # Reset module state
144+ db ._migrations_completed = False
145+ db ._engine = None
146+ db ._session_maker = None
147+
148+ # Call with ensure_migrations=False should skip migrations
149+ engine , session_maker = await db .get_or_create_db (
150+ app_config .database_path , ensure_migrations = False
151+ )
152+
153+ # Verify we got valid objects
154+ assert engine is not None
155+ assert session_maker is not None
156+
157+ # Verify migrations were NOT called
158+ mock_alembic_command .upgrade .assert_not_called ()
159+ mock_search_repository .init_search_index .assert_not_called ()
160+
161+
162+ @pytest .mark .asyncio
163+ async def test_multiple_get_or_create_db_calls_deduplicated (
164+ app_config , mock_alembic_config , mock_alembic_command , mock_search_repository
165+ ):
166+ """Test that multiple get_or_create_db calls only run migrations once."""
167+ # Reset module state
168+ db ._migrations_completed = False
169+ db ._engine = None
170+ db ._session_maker = None
171+
172+ # First call should create engine and run migrations
173+ await db .get_or_create_db (app_config .database_path , app_config = app_config )
174+
175+ # Verify migrations were called
176+ mock_alembic_command .upgrade .assert_called_once_with (mock_alembic_config , "head" )
177+ mock_search_repository .init_search_index .assert_called_once ()
178+
179+ # Reset mocks for subsequent calls
180+ mock_alembic_command .reset_mock ()
181+ mock_search_repository .reset_mock ()
182+
183+ # Subsequent calls should not run migrations again
184+ await db .get_or_create_db (app_config .database_path , app_config = app_config )
185+ await db .get_or_create_db (app_config .database_path , app_config = app_config )
186+
187+ # Verify migrations were NOT called again
188+ mock_alembic_command .upgrade .assert_not_called ()
189+ mock_search_repository .init_search_index .assert_not_called ()
0 commit comments