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