diff --git a/tests/test_app.py b/tests/test_app.py index 7e794e6a..f697afdf 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -602,3 +602,74 @@ async def test_rebuild_empty_database(tmp_path, monkeypatch): calls = [str(c) for c in mock_print.call_args_list] assert any("No documents found" in c for c in calls) + + +def test_migrate_with_pending_migrations(tmp_path): + """Test migrate method when migrations are applied.""" + from haiku.rag.store.engine import Store + + db_path = tmp_path / "test.lancedb" + store = Store(db_path, create=True) + store.close() + + app = HaikuRAGApp(db_path=db_path) + + with patch("haiku.rag.store.engine.Store") as mock_store_class: + mock_store = MagicMock() + mock_store.migrate.return_value = ["Migration 1", "Migration 2"] + mock_store_class.return_value = mock_store + + result = app.migrate() + + mock_store_class.assert_called_once_with( + db_path, + config=app.config, + skip_validation=True, + skip_migration_check=True, + ) + mock_store.migrate.assert_called_once() + mock_store.close.assert_called_once() + assert result == ["Migration 1", "Migration 2"] + + +def test_migrate_no_pending_migrations(tmp_path): + """Test migrate method when no migrations are pending.""" + from haiku.rag.store.engine import Store + + db_path = tmp_path / "test.lancedb" + store = Store(db_path, create=True) + store.close() + + app = HaikuRAGApp(db_path=db_path) + + with patch("haiku.rag.store.engine.Store") as mock_store_class: + mock_store = MagicMock() + mock_store.migrate.return_value = [] + mock_store_class.return_value = mock_store + + result = app.migrate() + + mock_store.migrate.assert_called_once() + mock_store.close.assert_called_once() + assert result == [] + + +def test_migrate_closes_store_on_exception(tmp_path): + """Test migrate method closes store even if migration fails.""" + from haiku.rag.store.engine import Store + + db_path = tmp_path / "test.lancedb" + store = Store(db_path, create=True) + store.close() + + app = HaikuRAGApp(db_path=db_path) + + with patch("haiku.rag.store.engine.Store") as mock_store_class: + mock_store = MagicMock() + mock_store.migrate.side_effect = Exception("Migration error") + mock_store_class.return_value = mock_store + + with pytest.raises(Exception, match="Migration error"): + app.migrate() + + mock_store.close.assert_called_once() diff --git a/tests/test_cli.py b/tests/test_cli.py index f3f8e266..01f02582 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3,6 +3,8 @@ from unittest.mock import AsyncMock, MagicMock, patch from typer.testing import CliRunner from haiku.rag.cli import _cli as cli +from haiku.rag.cli import cli as cli_wrapper +from haiku.rag.store.exceptions import MigrationRequiredError runner = CliRunner() @@ -386,3 +388,63 @@ def test_add_document_src_directory(tmp_path): mock_app_instance.add_document_from_source.assert_called_once() call_args = mock_app_instance.add_document_from_source.call_args assert call_args[1]["source"] == str(test_dir) + + +def test_migrate_with_applied_migrations(): + """Test migrate command when migrations are applied.""" + with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: + mock_app_instance = MagicMock() + mock_app_instance.migrate.return_value = [ + "Add full-text search index", + "Add metadata column", + ] + mock_app.return_value = mock_app_instance + + result = runner.invoke(cli, ["migrate"]) + + assert result.exit_code == 0 + mock_app_instance.migrate.assert_called_once() + assert "Applied 2 migration(s)" in result.output + assert "Add full-text search index" in result.output + assert "Add metadata column" in result.output + assert "Migration completed successfully" in result.output + + +def test_migrate_no_pending_migrations(): + """Test migrate command when no migrations are pending.""" + with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: + mock_app_instance = MagicMock() + mock_app_instance.migrate.return_value = [] + mock_app.return_value = mock_app_instance + + result = runner.invoke(cli, ["migrate"]) + + assert result.exit_code == 0 + mock_app_instance.migrate.assert_called_once() + assert "No migrations pending" in result.output + assert "Database is up to date" in result.output + + +def test_migrate_failure(): + """Test migrate command when migration fails.""" + with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: + mock_app_instance = MagicMock() + mock_app_instance.migrate.side_effect = Exception("Migration failed") + mock_app.return_value = mock_app_instance + + result = runner.invoke(cli, ["migrate"]) + + assert result.exit_code == 1 + assert "Migration failed" in result.output + + +def test_cli_wrapper_catches_migration_required_error(): + """Test that cli() wrapper catches MigrationRequiredError and exits with code 1.""" + with patch("haiku.rag.cli._cli") as mock_cli: + mock_cli.side_effect = MigrationRequiredError( + "Database requires migration. Run 'haiku-rag migrate' to upgrade." + ) + + with patch("sys.exit") as mock_exit: + cli_wrapper() + mock_exit.assert_called_once_with(1)