diff --git a/genesis_core/tests/functional/conftest.py b/genesis_core/tests/functional/conftest.py index 08a82a9c..aebbce49 100644 --- a/genesis_core/tests/functional/conftest.py +++ b/genesis_core/tests/functional/conftest.py @@ -915,3 +915,72 @@ def setup_db_for_worker(worker_id): engine.close_connection(conn) del engine engines.engine_factory.destroy_engine() + + + +@pytest.fixture(scope="session") +def main_db_url() -> str: + return consts.get_database_uri() + + +@pytest.fixture(scope="session") +def test_db_name(worker_id: tp.Any) -> str: + return "test" if not worker_id else f"test_{worker_id}" + + +@pytest.fixture(scope="session") +def test_db_url(main_db_url: str, test_db_name: str) -> str: + parsed_db_url = urlparse(main_db_url) + return parsed_db_url._replace(path=test_db_name).geturl() + + +@pytest.fixture(scope="session") +def test_db( + main_db_url: str, + test_db_url: str, + test_db_name: str, +) -> tp.Iterable[ + test_utils.TestDBManager +]: + class MainDBManager(test_utils.TestDBManager): + manager_config = test_utils.TestDBManagerConfig( + database_url=main_db_url, + create_db=test_db_name, + engine_alias="test_main", + ) + + class TestDBManager(test_utils.TestDBManager): + manager_config = test_utils.TestDBManagerConfig( + database_url=test_db_url, + engine_alias="test", + ) + + with MainDBManager() as main_db_manager: + with main_db_manager.db(): + with TestDBManager() as test_db_manager: + yield test_db_manager + + +@pytest.fixture(scope="session") +def test_migrations_manager( + test_db: test_utils.TestDBManager, +) -> tp.Iterable[None]: + class TestMigrationManager(test_utils.TestMigrationManager): + migration_config = test_utils.TestMigrationManagerConfig( + first_migration=FIRST_MIGRATION, + ) + + with TestMigrationManager(db_manager=test_db) as migration_manager: + with migration_manager.migrations(): + yield + + +@pytest.fixture() +def test_session( + test_db: test_utils.TestDBManager, + test_migrations_manager: None, +) -> tp.Iterable[ + test_utils.AbstractSession +]: + with test_db.session() as session: + yield session \ No newline at end of file diff --git a/genesis_core/tests/functional/dm/__init__.py b/genesis_core/tests/functional/dm/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/genesis_core/tests/functional/dm/test_machine_pool.py b/genesis_core/tests/functional/dm/test_machine_pool.py new file mode 100644 index 00000000..1df287c6 --- /dev/null +++ b/genesis_core/tests/functional/dm/test_machine_pool.py @@ -0,0 +1,119 @@ +# Copyright 2026 Genesis Corporation. +# +# All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may +# not use this file except in compliance with the License. You may obtain +# a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations +# under the License. + +import typing as tp + +import pytest + +from restalchemy.storage import exceptions + +from genesis_core.compute.dm.models import MachinePool +from genesis_core.tests.functional import utils as test_utils + + +DictStrAny = tp.Dict[str, tp.Any] + + +class DriverSpecWithException(tp.TypedDict): + driver_spec: tp.Optional[ + DictStrAny + ] + + exception: tp.Optional[ + tp.Type[Exception] + ] + + +DEFAULT_MACHINE_POOL_CONNECTION_URI = "qemu://system" +DEFAULT_DRIVER_SPEC_WITH_EXCEPTION = DriverSpecWithException( + driver_spec={"connection_uri": DEFAULT_MACHINE_POOL_CONNECTION_URI}, + exception=None, +) + + +class PoolFactory(tp.Protocol): + def __call__(self, *, driver_spec: DictStrAny) -> DictStrAny: + ... + + +@pytest.mark.parametrize( + "driver_specs_with_exceptions", + [ + pytest.param( + [ + DEFAULT_DRIVER_SPEC_WITH_EXCEPTION, + ], + id="single-insert" + ), + pytest.param( + [ + DriverSpecWithException(driver_spec=None, exception=None), + ], + id="none-driver-spec", + ), + pytest.param( + [ + DriverSpecWithException(driver_spec={}, exception=None), + ], + id="empty-dict-driver-spec", + ), + pytest.param( + [ + DriverSpecWithException( + driver_spec={"connection_uri": None}, + exception=None, + ), + ], + id="empty-connection-uri-driver-spec", + ), + pytest.param( + [ + DriverSpecWithException(driver_spec=None, exception=None), + DriverSpecWithException(driver_spec=None, exception=None), + ], + id="allow-connection-uri-null-duplicate", + ), + pytest.param( + [ + DEFAULT_DRIVER_SPEC_WITH_EXCEPTION, + DriverSpecWithException( + driver_spec={ + "connection_uri": DEFAULT_MACHINE_POOL_CONNECTION_URI, + }, + exception=exceptions.ConflictRecords, + ) + ], + id="disallow-duplicate-connection-uri", + ), + ], +) +def test_connection_uri_idx( + driver_specs_with_exceptions: tp.List[DriverSpecWithException], + test_session: test_utils.AbstractSession, + pool_factory: PoolFactory, +): + for param in driver_specs_with_exceptions: + machine_pool = MachinePool.restore_from_simple_view( + **pool_factory( + driver_spec=param["driver_spec"], + ) + ) + + if param["exception"] is None: + machine_pool.insert(session=test_session) + else: + with pytest.raises(param["exception"]): + machine_pool.insert(session=test_session) \ No newline at end of file diff --git a/genesis_core/tests/functional/utils.py b/genesis_core/tests/functional/utils.py index 3e969e21..b6dfa78e 100644 --- a/genesis_core/tests/functional/utils.py +++ b/genesis_core/tests/functional/utils.py @@ -15,17 +15,21 @@ # under the License. import os -import pathlib +import dataclasses import socket import contextlib +import pathlib from urllib import parse import typing as tp +from types import TracebackType from gcl_sdk import migrations as sdk_migrations -from restalchemy.storage.sql import migrations -from restalchemy.tests.functional import db_utils as ra_db_utils +from restalchemy.storage.sql import migrations, engines +from restalchemy.tests.functional import db_utils as ra_db_utils, consts as ra_consts from restalchemy.tests.functional.restapi.ra_based.microservice import service +from unittest.mock import patch + ENDPOINT_TEMPLATE = "http://127.0.0.1:%s/" @@ -170,3 +174,259 @@ def teardown_method(self) -> None: # Rollback migrations self._migration.rollback_migration(self.__FIRST_MIGRATION__) self._sdk_migration.rollback_migration(sdk_migrations.INIT_MIGRATION_FILENAME) + + +@contextlib.contextmanager +def patch_engines_default( + db_manager: "TestDBManager", +) -> tp.Iterable[None]: + default_name = engines.DEFAULT_NAME + default_get_engine = engines.EngineFactory.get_engine + + def _patched_get_engine( + factory: engines.EngineFactory, + name: str = default_name, + ) -> engines.AbstractEngine: + if name == default_name: + return db_manager.engine + else: + return default_get_engine(factory, name=name) + + with \ + patch( + "restalchemy.storage.sql.engines.DEFAULT_NAME", + db_manager.manager_config.engine_alias, + ), \ + patch.object( + engines.EngineFactory, + "get_engine", + _patched_get_engine, + ): + yield + + +OptionalStr = tp.Optional[str] + + +class AbstractCursor(tp.Protocol): + def execute(self, statement: str) -> tp.Any: + ... + + +class AbstractConnection(tp.Protocol): + autocommit: bool + + @contextlib.contextmanager + def cursor(self) -> tp.Iterable[AbstractCursor]: + ... + + def rollback(self) -> None: + ... + + +class AbstractSession(tp.Protocol): + def execute(self, statement, values=None) -> None: + ... + + def commit(self) -> None: + ... + + def rollback(self) -> None: + ... + + def close(self) -> None: + ... + + +@dataclasses.dataclass() +class TestDBManagerConfig: + database_url: OptionalStr = None + engine_alias: OptionalStr = None + create_db: OptionalStr = None + + def __post_init__(self) -> None: + self.database_url = self.database_url or ra_consts.get_database_uri() + self.engine_alias = self.engine_alias or engines.DEFAULT_NAME + + +class TestDBManager: + _Self = tp.TypeVar("_Self", bound="TestDBManager") + + manager_config: tp.ClassVar[TestDBManagerConfig] = TestDBManagerConfig() + + _engine: engines.AbstractEngine + + @property + def engine(self) -> engines.AbstractEngine: + return self._engine + + def setup(self) -> _Self: + engine_alias = self.manager_config.engine_alias + + engines.engine_factory.configure_factory( + db_url=self.manager_config.database_url, + name=engine_alias, + ) + self._engine = engines.engine_factory.get_engine(engine_alias) + + return self + + def teardown(self) -> None: + if isinstance(self._engine, engines.PgSQLEngine): + self._engine._pool.close() + + engines.engine_factory.destroy_engine( + name=self.manager_config.engine_alias, + ) + + def __enter__(self) -> _Self: + return self.setup() + + def __exit__( + self, + exc_type: tp.Type[Exception], + exc_val: Exception, + exc_tb: TracebackType + ) -> None: + self.teardown() + + @contextlib.contextmanager + def session(self) -> tp.Iterable[AbstractSession]: + session: AbstractSession = self._engine.get_session() + + try: + with patch_engines_default(db_manager=self): + yield session + finally: + session.rollback() + session.close() + + @contextlib.contextmanager + def connection( + self, + autocommit: bool = False, + ) -> tp.Iterable[AbstractConnection]: + connection: AbstractConnection = self._engine.get_connection() + + try: + connection.autocommit = autocommit + yield connection + finally: + if not autocommit: + connection.rollback() + + self._engine.close_connection(connection) + + def create_db(self) -> None: + create_db = self.manager_config.create_db + if create_db is None: + return + + with self.connection(autocommit=True) as connection: + with connection.cursor() as cursor: + cursor.execute(f"CREATE DATABASE \"{create_db}\"") + + def drop_db(self) -> None: + create_db = self.manager_config.create_db + if create_db is None: + return + + with self.connection(autocommit=True) as connection: + with connection.cursor() as cursor: + cursor.execute(f"DROP DATABASE \"{create_db}\"") + @contextlib.contextmanager + def db(self) -> tp.Iterable[_Self]: + try: + self.create_db() + yield self + + finally: + self.drop_db() + + +@dataclasses.dataclass() +class TestMigrationManagerConfig: + migrations_path: OptionalStr = None + first_migration: OptionalStr = None + last_migration: OptionalStr = None + + def __post_init__(self) -> None: + self.migrations_path = ( + self.migrations_path + or ( + str( + pathlib.Path(__file__) + .parent + .joinpath("../../../migrations/") + .resolve() + ) + ) + ) + + +class TestMigrationManager: + _Self = tp.TypeVar("_Self", bound="TestMigrationManager") + + migration_config: tp.ClassVar[ + TestMigrationManagerConfig + ] = TestMigrationManagerConfig() + + _migration_engine: migrations.MigrationEngine + + def __init__( + self, + db_manager: TestDBManager, + ) -> None: + self._db_manager = db_manager + + @property + def db_manager(self) -> TestDBManager: + return self._db_manager + + def setup(self) -> _Self: + self._migration_engine = migrations.MigrationEngine( + migrations_path=self.migration_config.migrations_path, + ) + + return self + + def __enter__(self) -> _Self: + return self.setup() + + def __exit__( + self, + exc_type: tp.Type[Exception], + exc_val: Exception, + exc_tb: TracebackType, + ) -> None: + return + + def apply_migrations(self) -> _Self: + last_migration = ( + self.migration_config.last_migration + or self._migration_engine.get_latest_migration() + ) + + with patch_engines_default( + db_manager=self.db_manager, + ): + self._migration_engine.apply_migration( + migration_name=last_migration, + ) + + return self + + def rollback_migrations(self) -> None: + with patch_engines_default( + db_manager=self.db_manager, + ): + self._migration_engine.rollback_migration( + migration_name=self.migration_config.first_migration, + ) + + @contextlib.contextmanager + def migrations(self) -> tp.Iterable[_Self]: + try: + yield self.apply_migrations() + finally: + self.rollback_migrations() \ No newline at end of file diff --git a/migrations/0055-unique-machine-pool-connection-uri-d16cdd.py b/migrations/0055-unique-machine-pool-connection-uri-d16cdd.py new file mode 100644 index 00000000..98f0624f --- /dev/null +++ b/migrations/0055-unique-machine-pool-connection-uri-d16cdd.py @@ -0,0 +1,71 @@ +# Copyright 2016 Eugene Frolov +# +# All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may +# not use this file except in compliance with the License. You may obtain +# a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations +# under the License. + +from restalchemy.storage.sql import migrations + + +class MigrationStep(migrations.AbstractMigrationStep): + + def __init__(self): + self._depends = ["0054-iam-idp-callback-kind-3a6c1b.py"] + + @property + def migration_id(self): + return "d16cddcf-b117-4a87-9ad7-8142925c23db" + + @property + def is_manual(self): + return False + + def upgrade(self, session): + expressions = [ + """ + CREATE OR REPLACE FUNCTION to_jsonb_safe(t TEXT) RETURNS jsonb AS $$ + BEGIN + RETURN t::jsonb; + EXCEPTION + WHEN invalid_text_representation THEN + RETURN NULL; + END; + $$ LANGUAGE plpgsql IMMUTABLE + """, + """ + CREATE UNIQUE INDEX IF NOT EXISTS + "machine_pools_connection_uri_unique_idx" + ON "machine_pools" ((to_jsonb_safe("driver_spec")->>'connection_uri')) + NULLS DISTINCT + """, + ] + + for expression in expressions: + session.execute(expression) + + def downgrade(self, session): + expressions = [ + """ + DROP INDEX IF EXISTS + "machine_pools_connection_uri_unique_idx" + """, + """ + DROP FUNCTION IF EXISTS to_jsonb_safe + """, + ] + + for expression in expressions: + session.execute(expression) + + +migration_step = MigrationStep() \ No newline at end of file