diff --git a/tests/aa_nodes/test_graph_execution_state.py b/tests/aa_nodes/test_graph_execution_state.py index d2850a12be..b5e5bde229 100644 --- a/tests/aa_nodes/test_graph_execution_state.py +++ b/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 diff --git a/tests/aa_nodes/test_invoker.py b/tests/aa_nodes/test_invoker.py index cc4d409b0f..da450a9fd2 100644 --- a/tests/aa_nodes/test_invoker.py +++ b/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 diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index eec6e033e5..dcf9c0a0f3 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/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 diff --git a/tests/app/services/model_records/test_model_records_sql.py b/tests/app/services/model_records/test_model_records_sql.py index 69cf753bb4..3df42f3ee1 100644 --- a/tests/app/services/model_records/test_model_records_sql.py +++ b/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) diff --git a/tests/test_sqlite_migrator.py b/tests/test_sqlite_migrator.py index 05beae799e..e1ca6ad5af 100644 --- a/tests/test_sqlite_migrator.py +++ b/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, )