From 3c7ad816d91d16c826d3771e338acac70f41ebe2 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Tue, 28 Jul 2026 12:29:02 +0900 Subject: [PATCH] test: use SQLite sessions in commands (#39080) --- .../test_legacy_model_type_migration.py | 322 +++++++++--------- 1 file changed, 161 insertions(+), 161 deletions(-) diff --git a/api/tests/unit_tests/commands/test_legacy_model_type_migration.py b/api/tests/unit_tests/commands/test_legacy_model_type_migration.py index 9eb73b82516..c5402df700e 100644 --- a/api/tests/unit_tests/commands/test_legacy_model_type_migration.py +++ b/api/tests/unit_tests/commands/test_legacy_model_type_migration.py @@ -6,6 +6,7 @@ import json import os import threading import time +from collections.abc import Iterator from datetime import datetime, timedelta from pathlib import Path from types import SimpleNamespace @@ -15,9 +16,12 @@ import pytest import sqlalchemy as sa from click.testing import CliRunner from sqlalchemy.exc import OperationalError +from sqlalchemy.orm import Session, SessionTransaction, sessionmaker from graphon.model_runtime.entities.model_entities import ModelType +from models import Dataset, DatasetPermission, DatasetPermissionEnum from models.account import Tenant +from models.base import TypeBase from models.enums import CredentialSourceType from models.provider import ProviderModel from tests.helpers.legacy_model_type_migration import ( @@ -59,6 +63,40 @@ def command_module(): ) +@pytest.fixture +def rbac_session(sqlite_engine: sa.Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]: + """Bind RBAC command reads to persisted SQLite dataset rows.""" + + TypeBase.metadata.create_all( + sqlite_engine, + tables=[Dataset.__table__, DatasetPermission.__table__], + ) + factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + monkeypatch.setattr("commands.rbac.session_factory.create_session", factory) + with factory() as session: + yield session + + +def _persist_dataset( + session: Session, + *, + dataset_id: str = "dataset-1", + tenant_id: str = "tenant-1", + permission: DatasetPermissionEnum = DatasetPermissionEnum.ONLY_ME, + created_by: str = "creator-account-1", +) -> Dataset: + dataset = Dataset( + id=dataset_id, + tenant_id=tenant_id, + name=f"Dataset {dataset_id}", + permission=permission, + created_by=created_by, + ) + session.add(dataset) + session.commit() + return dataset + + def _parse_json_lines(output: io.StringIO) -> list[dict[str, object]]: return [json.loads(line) for line in output.getvalue().splitlines() if line.strip()] @@ -363,56 +401,35 @@ def test_dataset_permission_rbac_migration_maps_legacy_permissions_to_enum_scope def test_dataset_permission_rbac_migration_uses_dataset_creator_as_operator( command_module, + rbac_session: Session, monkeypatch: pytest.MonkeyPatch, ) -> None: rbac_module = importlib.import_module("commands.rbac") - dataset_row = SimpleNamespace( - id="dataset-1", - tenant_id="tenant-1", - permission="only_me", - created_by="creator-account-1", - ) - execute_results = [[dataset_row], [], []] + _persist_dataset(rbac_session) calls: list[dict[str, object]] = [] - session_closed = False - - class FakeExecuteResult: - def __init__(self, rows: list[object]) -> None: - self._rows = rows - - def all(self) -> list[object]: - return self._rows - - class FakeSession: - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, traceback) -> None: - nonlocal session_closed - session_closed = True - pass - - def execute(self, stmt): - return FakeExecuteResult(execute_results.pop(0)) - - class FakeSessionFactory: - @staticmethod - def create_session() -> FakeSession: - return FakeSession() + read_transaction_ended = False def fake_replace_whitelist(**kwargs): - assert session_closed is True + assert read_transaction_ended is True calls.append(kwargs) - monkeypatch.setattr(rbac_module, "session_factory", FakeSessionFactory) - monkeypatch.setattr(rbac_module.RBACService.DatasetAccess, "replace_whitelist", fake_replace_whitelist) + def _record_transaction_end(session: Session, transaction: object) -> None: + nonlocal read_transaction_ended + del transaction + if session.get_bind() is rbac_session.get_bind(): + read_transaction_ended = True - command_module.migrate_dataset_permissions_to_rbac.callback( - tenant_id=None, - dataset_id=None, - batch_size=500, - dry_run=False, - ) + sa.event.listen(Session, "after_transaction_end", _record_transaction_end) + monkeypatch.setattr(rbac_module.RBACService.DatasetAccess, "replace_whitelist", fake_replace_whitelist) + try: + command_module.migrate_dataset_permissions_to_rbac.callback( + tenant_id=None, + dataset_id=None, + batch_size=500, + dry_run=False, + ) + finally: + sa.event.remove(Session, "after_transaction_end", _record_transaction_end) assert calls[0]["tenant_id"] == "tenant-1" assert calls[0]["account_id"] == "creator-account-1" @@ -422,41 +439,19 @@ def test_dataset_permission_rbac_migration_uses_dataset_creator_as_operator( def test_dataset_permission_rbac_migration_dry_run_outputs_structured_proposed_changes( command_module, + rbac_session: Session, monkeypatch: pytest.MonkeyPatch, ) -> None: rbac_module = importlib.import_module("commands.rbac") - dataset_row = SimpleNamespace( - id="dataset-1", - tenant_id="tenant-1", - permission="partial_members", - created_by="creator-account-1", + dataset = _persist_dataset(rbac_session, permission=DatasetPermissionEnum.PARTIAL_TEAM) + rbac_session.add( + DatasetPermission( + dataset_id=dataset.id, + account_id="member-account-1", + tenant_id=dataset.tenant_id, + ) ) - permission_row = SimpleNamespace(dataset_id="dataset-1", account_id="member-account-1") - execute_results = [[dataset_row], [permission_row], []] - - class FakeExecuteResult: - def __init__(self, rows: list[object]) -> None: - self._rows = rows - - def all(self) -> list[object]: - return self._rows - - class FakeSession: - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, traceback) -> None: - pass - - def execute(self, stmt): - return FakeExecuteResult(execute_results.pop(0)) - - class FakeSessionFactory: - @staticmethod - def create_session() -> FakeSession: - return FakeSession() - - monkeypatch.setattr(rbac_module, "session_factory", FakeSessionFactory) + rbac_session.commit() monkeypatch.setattr( rbac_module.RBACService.DatasetAccess, "replace_whitelist", @@ -1306,50 +1301,36 @@ def test_provider_models_processing_uses_same_plan_locking_and_transaction_entry begin_calls: list[str] = [] configure_calls: list[str] = [] - class _FakeBeginContext: - def __init__(self, phase: str) -> None: - self._phase = phase + def _record_begin(session: Session, transaction: SessionTransaction) -> None: + if session.get_bind() is sqlite_engine and transaction.parent is None: + begin_calls.append(current_phase["name"]) - def __enter__(self) -> None: - begin_calls.append(self._phase) - - def __exit__(self, exc_type, exc, tb) -> bool: - return False - - class _FakeSession: - def __init__(self, phase: str) -> None: - self._phase = phase - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb) -> bool: - return False - - def begin(self) -> _FakeBeginContext: - return _FakeBeginContext(self._phase) - - def _fake_session_factory(engine: sa.Engine) -> _FakeSession: - return _FakeSession(current_phase["name"]) - - def _fake_build_plan(self, session, candidate, *, lock_rows: bool): + def _fake_build_plan(self, session: Session, candidate, *, lock_rows: bool): + assert session.get_bind() is sqlite_engine lock_rows_seen.append((current_phase["name"], lock_rows)) - return SimpleNamespace(group_row_ids=[str(candidate.row.id)], winner=None, loser_rows=[]) + return migration_module._ProviderModelGroupPlan( + group_row_ids=[str(candidate.row.id)], + winner=None, + loser_rows=[], + ) def _fake_emit_plan(self, plan, *, session, tx_id: str, business_key: dict[str, object]) -> None: return None - def _fake_configure(self, session) -> None: + def _fake_configure(self, session: Session) -> None: + assert session.get_bind() is sqlite_engine configure_calls.append(current_phase["name"]) - monkeypatch.setattr(migration_module, "_session_factory", _fake_session_factory) monkeypatch.setattr(migration_module.Migration, "_build_provider_model_group_plan", _fake_build_plan) monkeypatch.setattr(migration_module.Migration, "_emit_provider_model_group_plan", _fake_emit_plan) monkeypatch.setattr(migration_module.Migration, "_configure_lock_timeout", _fake_configure) - - dry_migration._process_provider_model_group(candidate, business_key) - current_phase["name"] = "apply" - apply_migration._process_provider_model_group(candidate, business_key) + sa.event.listen(Session, "after_transaction_create", _record_begin) + try: + dry_migration._process_provider_model_group(candidate, business_key) + current_phase["name"] = "apply" + apply_migration._process_provider_model_group(candidate, business_key) + finally: + sa.event.remove(Session, "after_transaction_create", _record_begin) assert [phase for phase, _ in lock_rows_seen] == ["dry", "apply"] assert lock_rows_seen[0][1] == lock_rows_seen[1][1] @@ -1392,6 +1373,22 @@ def test_process_load_balancing_model_config_row_logs_stacktrace_for_lock_timeou sqlite_engine: sa.Engine, monkeypatch: pytest.MonkeyPatch, ) -> None: + create_minimal_legacy_model_type_schema(sqlite_engine) + created_at = datetime(2025, 1, 1, 12, 0, 0) + _insert_load_balancing_model_config( + sqlite_engine, + row_id="40000000-0000-0000-0000-000000000001", + tenant_id="tenant-1", + provider_name="openai", + model_name="gpt-4o-mini", + model_type="text-generation", + name="credential", + encrypted_config="{}", + credential_id="50000000-0000-0000-0000-000000000001", + enabled=True, + created_at=created_at, + updated_at=created_at, + ) output = io.StringIO() migration = migration_module.Migration( tenant_id="tenant-1", @@ -1401,37 +1398,18 @@ def test_process_load_balancing_model_config_row_logs_stacktrace_for_lock_timeou model_types=(ModelType.LLM,), orm_models=(migration_module.LoadBalancingModelConfig,), ) - candidate = migration_module._RowWithRawModelType( - row=SimpleNamespace(id="lb-row-1"), - raw_model_type="text-generation", - canonical_model_type=ModelType.LLM, - ) + candidate = migration._load_load_balancing_model_config_candidates(None)[0] lock_timeout_exc = OperationalError("SELECT 1", {}, SimpleNamespace(pgcode="55P03")) + transaction_begins = 0 - class _FakeBeginContext: - def __enter__(self) -> None: - return None - - def __exit__(self, exc_type, exc, tb) -> bool: - return False - - class _FakeSession: - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb) -> bool: - return False - - def begin(self) -> _FakeBeginContext: - return _FakeBeginContext() - - def _fake_session_factory(engine: sa.Engine) -> _FakeSession: - return _FakeSession() + def _record_begin(session: Session, transaction: SessionTransaction) -> None: + nonlocal transaction_begins + if session.get_bind() is sqlite_engine and transaction.parent is None: + transaction_begins += 1 def _fake_reload(self, session, original_candidate, *, lock_rows: bool): raise lock_timeout_exc - monkeypatch.setattr(migration_module, "_session_factory", _fake_session_factory) monkeypatch.setattr(migration_module.Migration, "_configure_lock_timeout", lambda self, session: None) monkeypatch.setattr( migration_module.Migration, @@ -1439,17 +1417,22 @@ def test_process_load_balancing_model_config_row_logs_stacktrace_for_lock_timeou _fake_reload, ) - migration._process_load_balancing_model_config_row(candidate) + sa.event.listen(Session, "after_transaction_create", _record_begin) + try: + migration._process_load_balancing_model_config_row(candidate) + finally: + sa.event.remove(Session, "after_transaction_create", _record_begin) lines = _parse_json_lines(output) assert len(lines) == 1 assert lines[0]["event"] == "lock_timeout_skipped" attrs = cast(dict[str, object], lines[0]["attrs"]) assert attrs["table_name"] == "load_balancing_model_configs" - assert attrs["id"] == "lb-row-1" + assert attrs["id"] == str(candidate.row.id) assert attrs["error"] == str(lock_timeout_exc) assert isinstance(attrs["stacktrace"], str) assert "OperationalError" in attrs["stacktrace"] + assert transaction_begins == 1 def test_process_load_balancing_model_config_row_logs_update_after_sql_execution( @@ -1457,6 +1440,23 @@ def test_process_load_balancing_model_config_row_logs_update_after_sql_execution sqlite_engine: sa.Engine, monkeypatch: pytest.MonkeyPatch, ) -> None: + create_minimal_legacy_model_type_schema(sqlite_engine) + created_at = datetime(2025, 1, 1, 12, 0, 0) + row_id = "40000000-0000-0000-0000-000000000002" + _insert_load_balancing_model_config( + sqlite_engine, + row_id=row_id, + tenant_id="tenant-1", + provider_name="openai", + model_name="gpt-4o-mini", + model_type="text-generation", + name="credential", + encrypted_config="{}", + credential_id="50000000-0000-0000-0000-000000000002", + enabled=True, + created_at=created_at, + updated_at=created_at, + ) migration = migration_module.Migration( tenant_id="tenant-1", engine=sqlite_engine, @@ -1465,42 +1465,33 @@ def test_process_load_balancing_model_config_row_logs_update_after_sql_execution model_types=(ModelType.LLM,), orm_models=(migration_module.LoadBalancingModelConfig,), ) - candidate = migration_module._RowWithRawModelType( - row=SimpleNamespace(id="lb-row-1"), - raw_model_type="text-generation", - canonical_model_type=ModelType.LLM, - ) + candidate = migration._load_load_balancing_model_config_candidates(None)[0] action_log: list[str] = [] - class _FakeBeginContext: - def __enter__(self) -> None: + def _record_begin(session: Session, transaction: SessionTransaction) -> None: + if session.get_bind() is sqlite_engine and transaction.parent is None: action_log.append("begin") - def __exit__(self, exc_type, exc, tb) -> bool: - return False - - class _FakeSession: - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb) -> bool: - return False - - def begin(self) -> _FakeBeginContext: - return _FakeBeginContext() - - def execute(self, stmt) -> None: + def _record_sql( + connection: sa.Connection, + cursor: object, + statement: str, + parameters: object, + context: object, + executemany: bool, + ) -> None: + del connection, cursor, parameters, context, executemany + if statement.lstrip().upper().startswith("UPDATE"): action_log.append("sql_execute") - def _fake_session_factory(engine: sa.Engine) -> _FakeSession: - return _FakeSession() - def _fake_configure(self, session) -> None: action_log.append("configure_lock_timeout") - def _fake_reload(self, session, original_candidate, *, lock_rows: bool): + original_reload = migration_module.Migration._reload_load_balancing_model_config_candidate + + def _record_reload(self, session: Session, original_candidate, *, lock_rows: bool): action_log.append(f"reload_candidate:{lock_rows}") - return candidate + return original_reload(self, session, original_candidate, lock_rows=lock_rows) def _fake_log_row_updated(self, *args, **kwargs) -> None: action_log.append("log_row_updated") @@ -1508,12 +1499,11 @@ def test_process_load_balancing_model_config_row_logs_update_after_sql_execution def _fake_cache_cleanup(self, *, row_id: str, tx_id: str) -> None: action_log.append("cache_cleanup") - monkeypatch.setattr(migration_module, "_session_factory", _fake_session_factory) monkeypatch.setattr(migration_module.Migration, "_configure_lock_timeout", _fake_configure) monkeypatch.setattr( migration_module.Migration, "_reload_load_balancing_model_config_candidate", - _fake_reload, + _record_reload, ) monkeypatch.setattr(migration_module.Migration, "_log_row_updated", _fake_log_row_updated) monkeypatch.setattr( @@ -1522,7 +1512,13 @@ def test_process_load_balancing_model_config_row_logs_update_after_sql_execution _fake_cache_cleanup, ) - migration._process_load_balancing_model_config_row(candidate) + sa.event.listen(Session, "after_transaction_create", _record_begin) + sa.event.listen(sqlite_engine, "before_cursor_execute", _record_sql) + try: + migration._process_load_balancing_model_config_row(candidate) + finally: + sa.event.remove(sqlite_engine, "before_cursor_execute", _record_sql) + sa.event.remove(Session, "after_transaction_create", _record_begin) assert action_log == [ "begin", @@ -1532,6 +1528,10 @@ def test_process_load_balancing_model_config_row_logs_update_after_sql_execution "log_row_updated", "cache_cleanup", ] + with Session(sqlite_engine) as session: + persisted = session.get(migration_module.LoadBalancingModelConfig, row_id) + assert persisted is not None + assert persisted.model_type == ModelType.LLM def test_load_balancing_model_config_cache_delete_failure_logs_stacktrace(