fix(db): fix tests

tests/aa_nodes/test_graph_execution_state.py

@@ -29,10 +29,10 @@ from invokeai.app.services.shared.graph import (
     IterateInvocation,
     LibraryGraph,
 )
-from invokeai.app.services.shared.sqlite.migrations.migration_1 import migration_1
-from invokeai.app.services.shared.sqlite.migrations.migration_2 import migration_2
 from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
-from invokeai.app.services.shared.sqlite.sqlite_migrator import SQLiteMigrator
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_1 import migration_1
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_2 import migration_2
+from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SQLiteMigrator
 from invokeai.backend.util.logging import InvokeAILogger
 
 from .test_invoker import create_edge

tests/aa_nodes/test_invoker.py

@@ -4,7 +4,7 @@ from pathlib import Path
 import pytest
 
 from invokeai.app.services.config.config_default import InvokeAIAppConfig
-from invokeai.app.services.shared.sqlite.sqlite_migrator import SQLiteMigrator
+from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SQLiteMigrator
 from invokeai.backend.util.logging import InvokeAILogger
 
 # This import must happen before other invoke imports or test in other files(!!) break
@@ -26,9 +26,9 @@ from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.item_storage.item_storage_sqlite import SqliteItemStorage
 from invokeai.app.services.session_queue.session_queue_common import DEFAULT_QUEUE_ID
 from invokeai.app.services.shared.graph import Graph, GraphExecutionState, GraphInvocation, LibraryGraph
-from invokeai.app.services.shared.sqlite.migrations.migration_1 import migration_1
-from invokeai.app.services.shared.sqlite.migrations.migration_2 import migration_2
 from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_1 import migration_1
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_2 import migration_2
 
 
 @pytest.fixture

tests/app/services/model_install/test_model_install.py

@@ -19,6 +19,9 @@ from invokeai.app.services.model_install import (
 )
 from invokeai.app.services.model_records import ModelRecordServiceBase, ModelRecordServiceSQL, UnknownModelException
 from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_1 import migration_1
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_2 import migration_2
+from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SQLiteMigrator
 from invokeai.backend.model_manager.config import BaseModelType, ModelType
 from invokeai.backend.util.logging import InvokeAILogger
 
@@ -38,7 +41,20 @@ def app_config(datadir: Path) -> InvokeAIAppConfig:
 
 @pytest.fixture
 def store(app_config: InvokeAIAppConfig) -> ModelRecordServiceBase:
-    database = SqliteDatabase(app_config, InvokeAILogger.get_logger(config=app_config))
+    logger = InvokeAILogger.get_logger(config=app_config)
+    database = SqliteDatabase(app_config, logger)
+    migrator = SQLiteMigrator(
+        db_path=database.database if isinstance(database.database, Path) else None,
+        conn=database.conn,
+        lock=database.lock,
+        logger=logger,
+        log_sql=app_config.log_sql,
+    )
+    migrator.register_migration(migration_1)
+    migrator.register_migration(migration_2)
+    migrator.run_migrations()
+    # this test uses a file database, so we need to reinitialize it after migrations
+    database.reinitialize()
     store: ModelRecordServiceBase = ModelRecordServiceSQL(database)
     return store
 

tests/app/services/model_records/test_model_records_sql.py

@@ -4,6 +4,7 @@ Test the refactored model config classes.
 
 from hashlib import sha256
 from pathlib import Path
+from typing import Any
 
 import pytest
 
@@ -14,10 +15,10 @@ from invokeai.app.services.model_records import (
     ModelRecordServiceSQL,
     UnknownModelException,
 )
-from invokeai.app.services.shared.sqlite.migrations.migration_1 import migration_1
-from invokeai.app.services.shared.sqlite.migrations.migration_2 import migration_2
 from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
-from invokeai.app.services.shared.sqlite.sqlite_migrator import SQLiteMigrator
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_1 import migration_1
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_2 import migration_2
+from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SQLiteMigrator
 from invokeai.backend.model_manager.config import (
     BaseModelType,
     MainCheckpointConfig,
@@ -30,7 +31,7 @@ from invokeai.backend.util.logging import InvokeAILogger
 
 
 @pytest.fixture
-def store(datadir) -> ModelRecordServiceBase:
+def store(datadir: Any) -> ModelRecordServiceBase:
     config = InvokeAIAppConfig(root=datadir)
     logger = InvokeAILogger.get_logger(config=config)
     db = SqliteDatabase(config, logger)

tests/test_sqlite_migrator.py

@@ -9,12 +9,14 @@ import pytest
 from pydantic import ValidationError
 
 from invokeai.app.services.shared.sqlite.sqlite_common import sqlite_memory
-from invokeai.app.services.shared.sqlite.sqlite_migrator import (
+from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import (
     MigrateCallback,
     Migration,
     MigrationError,
     MigrationSet,
     MigrationVersionError,
+)
+from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import (
     SQLiteMigrator,
 )