From f15a51487812cc4e63560082031395660f59c98c Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Tue, 4 Aug 2026 00:44:54 +0900 Subject: [PATCH] test: use SQLite sessions in core test_provider_manager.py (#39071) Co-authored-by: Byron Wang --- .../unit_tests/core/test_provider_manager.py | 544 ++++++++++++------ 1 file changed, 362 insertions(+), 182 deletions(-) diff --git a/api/tests/unit_tests/core/test_provider_manager.py b/api/tests/unit_tests/core/test_provider_manager.py index 128eebdd5af..7c9547a6492 100644 --- a/api/tests/unit_tests/core/test_provider_manager.py +++ b/api/tests/unit_tests/core/test_provider_manager.py @@ -1,33 +1,74 @@ +from collections.abc import Iterator from types import SimpleNamespace from unittest.mock import MagicMock, Mock, PropertyMock, patch import pytest -from pytest_mock import MockerFixture +from sqlalchemy import Engine, event, select +from sqlalchemy.orm import Session, sessionmaker +from core import provider_manager as provider_manager_module from core.entities.provider_entities import ModelSettings from core.provider_manager import ProviderConfigurationCacheSource, ProviderManager from graphon.model_runtime.entities.common_entities import I18nObject from graphon.model_runtime.entities.model_entities import ModelType +from models.base import TypeBase from models.provider import ( LoadBalancingModelConfig, Provider, ProviderCredential, + ProviderModel, + ProviderModelCredential, ProviderModelSetting, ProviderType, TenantDefaultModel, + TenantPreferredModelProvider, ) from models.provider_ids import ModelProviderID -def _build_provider_manager(mocker: MockerFixture) -> ProviderManager: - return ProviderManager(model_runtime=mocker.Mock()) +def _build_provider_manager() -> ProviderManager: + return ProviderManager(model_runtime=Mock()) -def _build_session_context(session: Mock) -> MagicMock: - session_cm = MagicMock() - session_cm.__enter__.return_value = session - session_cm.__exit__.return_value = False - return session_cm +def _persist_model_configuration( + session: Session, + setting: ProviderModelSetting, + load_balancing_configs: list[LoadBalancingModelConfig], +) -> tuple[list[ProviderModelSetting], list[LoadBalancingModelConfig]]: + session.add_all([setting, *load_balancing_configs]) + session.commit() + session.expire_all() + settings = list( + session.scalars(select(ProviderModelSetting).where(ProviderModelSetting.tenant_id == setting.tenant_id)).all() + ) + configs = list( + session.scalars( + select(LoadBalancingModelConfig).where(LoadBalancingModelConfig.tenant_id == setting.tenant_id) + ).all() + ) + return settings, configs + + +@pytest.fixture +def provider_db(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]: + """Bind request-owned and provider-owned sessions to one isolated SQLite database.""" + models = ( + Provider, + ProviderCredential, + ProviderModel, + ProviderModelCredential, + ProviderModelSetting, + LoadBalancingModelConfig, + TenantDefaultModel, + TenantPreferredModelProvider, + ) + tables = [model.metadata.tables[model.__tablename__] for model in models] + TypeBase.metadata.create_all(sqlite_engine, tables=tables) + owned_session_factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + with owned_session_factory() as request_session: + monkeypatch.setattr(provider_manager_module.db, "session", request_session) + monkeypatch.setattr(provider_manager_module.session_factory, "create_session", owned_session_factory) + yield request_session class _FakeRedis: @@ -56,17 +97,6 @@ class _FakeRedis: self.expirations[key] = time -class _FakeScalarResult: - def __init__(self, values: list[object]) -> None: - self._values = values - - def __iter__(self): - return iter(self._values) - - def all(self) -> list[object]: - return self._values - - @pytest.fixture def mock_provider_entity(): mock_entity = Mock() @@ -86,7 +116,7 @@ def mock_provider_entity(): return mock_entity -def test__to_model_settings(mocker: MockerFixture, mock_provider_entity): +def test__to_model_settings(mock_provider_entity, provider_db: Session): # Mocking the inputs ps = ProviderModelSetting( tenant_id="tenant_id", @@ -122,12 +152,15 @@ def test__to_model_settings(mocker: MockerFixture, mock_provider_entity): ] load_balancing_model_configs[0].id = "id1" load_balancing_model_configs[1].id = "id2" + provider_model_settings, load_balancing_model_configs = _persist_model_configuration( + provider_db, ps, load_balancing_model_configs + ) with patch( "core.helper.model_provider_cache.ProviderCredentialsCache.get", return_value={"openai_api_key": "fake_key"}, ): - provider_manager = _build_provider_manager(mocker) + provider_manager = _build_provider_manager() # Running the method result = provider_manager._to_model_settings( @@ -147,7 +180,7 @@ def test__to_model_settings(mocker: MockerFixture, mock_provider_entity): assert result[0].load_balancing_configs[1].name == "first" -def test__to_model_settings_only_one_lb(mocker: MockerFixture, mock_provider_entity): +def test__to_model_settings_only_one_lb(mock_provider_entity, provider_db: Session): # Mocking the inputs ps = ProviderModelSetting( @@ -172,12 +205,15 @@ def test__to_model_settings_only_one_lb(mocker: MockerFixture, mock_provider_ent ) ] load_balancing_model_configs[0].id = "id1" + provider_model_settings, load_balancing_model_configs = _persist_model_configuration( + provider_db, ps, load_balancing_model_configs + ) with patch( "core.helper.model_provider_cache.ProviderCredentialsCache.get", return_value={"openai_api_key": "fake_key"}, ): - provider_manager = _build_provider_manager(mocker) + provider_manager = _build_provider_manager() # Running the method result = provider_manager._to_model_settings( @@ -195,7 +231,7 @@ def test__to_model_settings_only_one_lb(mocker: MockerFixture, mock_provider_ent assert len(result[0].load_balancing_configs) == 0 -def test__to_model_settings_lb_disabled(mocker: MockerFixture, mock_provider_entity): +def test__to_model_settings_lb_disabled(mock_provider_entity, provider_db: Session): # Mocking the inputs ps = ProviderModelSetting( tenant_id="tenant_id", @@ -229,12 +265,15 @@ def test__to_model_settings_lb_disabled(mocker: MockerFixture, mock_provider_ent ] load_balancing_model_configs[0].id = "id1" load_balancing_model_configs[1].id = "id2" + provider_model_settings, load_balancing_model_configs = _persist_model_configuration( + provider_db, ps, load_balancing_model_configs + ) with patch( "core.helper.model_provider_cache.ProviderCredentialsCache.get", return_value={"openai_api_key": "fake_key"}, ): - provider_manager = _build_provider_manager(mocker) + provider_manager = _build_provider_manager() # Running the method result = provider_manager._to_model_settings( @@ -252,19 +291,23 @@ def test__to_model_settings_lb_disabled(mocker: MockerFixture, mock_provider_ent assert len(result[0].load_balancing_configs) == 0 -def test_get_default_model_uses_first_available_active_model(mocker: MockerFixture): - mock_session = Mock() - mock_session.scalar.return_value = None - +def test_get_default_model_uses_first_available_active_model(provider_db: Session): + other_tenant_default = TenantDefaultModel( + tenant_id="other-tenant", + provider_name="anthropic", + model_name="claude", + model_type=ModelType.LLM, + ) + provider_db.add(other_tenant_default) + provider_db.commit() provider_configurations = Mock() provider_configurations.get_models.return_value = [ Mock(model="gpt-3.5-turbo", provider=Mock(provider="openai")), Mock(model="gpt-4", provider=Mock(provider="openai")), ] - manager = _build_provider_manager(mocker) + manager = _build_provider_manager() with ( - patch("core.provider_manager.db.session", mock_session), patch.object(manager, "get_configurations", return_value=provider_configurations), patch("core.provider_manager.ModelProviderFactory") as mock_factory_cls, ): @@ -281,22 +324,33 @@ def test_get_default_model_uses_first_available_active_model(mocker: MockerFixtu assert result.model == "gpt-3.5-turbo" assert result.provider.provider == "openai" provider_configurations.get_models.assert_called_once_with(model_type=ModelType.LLM, only_active=True) - mock_session.add.assert_called_once() - saved_default_model = mock_session.add.call_args.args[0] + saved_default_model = provider_db.scalar( + select(TenantDefaultModel).where(TenantDefaultModel.tenant_id == "tenant-id") + ) + assert saved_default_model is not None assert saved_default_model.model_name == "gpt-3.5-turbo" assert saved_default_model.provider_name == "openai" - mock_session.commit.assert_called_once() + assert ( + provider_db.scalar(select(TenantDefaultModel).where(TenantDefaultModel.tenant_id == "other-tenant")) + is other_tenant_default + ) -def test_get_default_model_returns_none_when_no_default_or_active_models(mocker: MockerFixture): - mock_session = Mock() - mock_session.scalar.return_value = None +def test_get_default_model_returns_none_when_no_default_or_active_models(provider_db: Session): + provider_db.add( + TenantDefaultModel( + tenant_id="other-tenant", + provider_name="anthropic", + model_name="claude", + model_type=ModelType.LLM, + ) + ) + provider_db.commit() provider_configurations = Mock() provider_configurations.get_models.return_value = [] - manager = _build_provider_manager(mocker) + manager = _build_provider_manager() with ( - patch("core.provider_manager.db.session", mock_session), patch.object(manager, "get_configurations", return_value=provider_configurations), patch("core.provider_manager.ModelProviderFactory") as mock_factory_cls, ): @@ -305,23 +359,24 @@ def test_get_default_model_returns_none_when_no_default_or_active_models(mocker: assert result is None provider_configurations.get_models.assert_called_once_with(model_type=ModelType.LLM, only_active=True) mock_factory_cls.assert_not_called() - mock_session.add.assert_not_called() - mock_session.commit.assert_not_called() + assert provider_db.scalar(select(TenantDefaultModel).where(TenantDefaultModel.tenant_id == "tenant-id")) is None + assert ( + provider_db.scalar(select(TenantDefaultModel).where(TenantDefaultModel.tenant_id == "other-tenant")) is not None + ) -def test_get_default_model_uses_injected_runtime_for_existing_default_record(mocker: MockerFixture): +def test_get_default_model_uses_injected_runtime_for_existing_default_record(provider_db: Session): existing_default_model = TenantDefaultModel( tenant_id="tenant-id", provider_name="openai", model_name="gpt-4", model_type=ModelType.LLM, ) - mock_session = Mock() - mock_session.scalar.return_value = existing_default_model - manager = _build_provider_manager(mocker) + provider_db.add(existing_default_model) + provider_db.commit() + manager = _build_provider_manager() with ( - patch("core.provider_manager.db.session", mock_session), patch("core.provider_manager.ModelProviderFactory") as mock_factory_cls, ): mock_factory_cls.return_value.get_provider_schema.return_value = Mock( @@ -339,11 +394,32 @@ def test_get_default_model_uses_injected_runtime_for_existing_default_record(moc assert result.provider.provider == "openai" -def test_get_configurations_uses_injected_runtime_and_adds_provider_aliases(mocker: MockerFixture): - manager = _build_provider_manager(mocker) - provider_records = {"openai": [SimpleNamespace(provider_name="openai")]} - provider_model_records = {"openai": [SimpleNamespace(provider_name="openai")]} - preferred_provider_records = {"openai": SimpleNamespace(preferred_provider_type="system")} +def test_get_configurations_uses_injected_runtime_and_adds_provider_aliases(provider_db: Session): + manager = _build_provider_manager() + provider = Provider( + tenant_id="tenant-id", + provider_name="openai", + provider_type=ProviderType.CUSTOM, + is_valid=True, + ) + provider_model = ProviderModel( + tenant_id="tenant-id", + provider_name="openai", + model_name="gpt-4", + model_type=ModelType.LLM, + is_valid=True, + ) + preferred_provider = TenantPreferredModelProvider( + tenant_id="tenant-id", + provider_name="openai", + preferred_provider_type=ProviderType.SYSTEM, + ) + provider_db.add_all([provider, provider_model, preferred_provider]) + provider_db.commit() + with patch("core.provider_manager.redis_client", _FakeRedis()): + provider_model_records = ProviderManager._get_all_provider_models("tenant-id") + preferred_provider_records = ProviderManager._get_all_preferred_model_providers("tenant-id") + provider_records = {"openai": [provider]} with ( patch.object(manager, "_get_all_providers", return_value=provider_records), @@ -380,18 +456,16 @@ def test_get_provider_names_returns_short_and_full_aliases(provider_name: str, e assert ProviderManager._get_provider_names(provider_name) == expected_provider_names -def test_get_provider_model_bundle_raises_for_unknown_provider(mocker: MockerFixture): - manager = _build_provider_manager(mocker) +def test_get_provider_model_bundle_raises_for_unknown_provider(): + manager = _build_provider_manager() with patch.object(manager, "get_configurations", return_value={}): with pytest.raises(ValueError, match="Provider openai does not exist."): manager.get_provider_model_bundle("tenant-id", "openai", ModelType.LLM) -def test_get_configurations_binds_manager_runtime_to_provider_configuration( - mocker: MockerFixture, mock_provider_entity -): - manager = _build_provider_manager(mocker) +def test_get_configurations_binds_manager_runtime_to_provider_configuration(mock_provider_entity): + manager = _build_provider_manager() provider_configuration = Mock() provider_factory = Mock() provider_factory.get_providers.return_value = [mock_provider_entity] @@ -418,8 +492,8 @@ def test_get_configurations_binds_manager_runtime_to_provider_configuration( provider_configuration.bind_model_runtime.assert_called_once_with(manager._model_runtime) -def test_get_configurations_reuses_cached_result_for_same_tenant(mocker: MockerFixture, mock_provider_entity): - manager = _build_provider_manager(mocker) +def test_get_configurations_reuses_cached_result_for_same_tenant(mock_provider_entity): + manager = _build_provider_manager() provider_configuration = Mock() provider_factory = Mock() provider_factory.get_providers.return_value = [mock_provider_entity] @@ -454,8 +528,8 @@ def test_get_configurations_reuses_cached_result_for_same_tenant(mocker: MockerF provider_configuration.bind_model_runtime.assert_called_once_with(manager._model_runtime) -def test_clear_configurations_cache_rebuilds_requested_tenant(mocker: MockerFixture, mock_provider_entity): - manager = _build_provider_manager(mocker) +def test_clear_configurations_cache_rebuilds_requested_tenant(mock_provider_entity): + manager = _build_provider_manager() provider_factory = Mock() provider_factory.get_providers.return_value = [mock_provider_entity] custom_configuration = SimpleNamespace(provider=None, models=[]) @@ -492,8 +566,8 @@ def test_clear_configurations_cache_rebuilds_requested_tenant(mocker: MockerFixt provider_configuration_second.bind_model_runtime.assert_called_once_with(manager._model_runtime) -def test_get_provider_model_bundle_returns_selected_model_type_instance(mocker: MockerFixture): - manager = _build_provider_manager(mocker) +def test_get_provider_model_bundle_returns_selected_model_type_instance(): + manager = _build_provider_manager() provider_configuration = Mock() model_type_instance = Mock() provider_configuration.get_model_type_instance.return_value = model_type_instance @@ -513,8 +587,8 @@ def test_get_provider_model_bundle_returns_selected_model_type_instance(mocker: assert result is expected_bundle -def test_get_first_provider_first_model_returns_none_when_no_models(mocker: MockerFixture): - manager = _build_provider_manager(mocker) +def test_get_first_provider_first_model_returns_none_when_no_models(): + manager = _build_provider_manager() provider_configurations = Mock() provider_configurations.get_models.return_value = [] @@ -525,8 +599,8 @@ def test_get_first_provider_first_model_returns_none_when_no_models(mocker: Mock provider_configurations.get_models.assert_called_once_with(model_type=ModelType.LLM, only_active=False) -def test_get_first_provider_first_model_returns_first_model_and_provider(mocker: MockerFixture): - manager = _build_provider_manager(mocker) +def test_get_first_provider_first_model_returns_first_model_and_provider(): + manager = _build_provider_manager() provider_configurations = Mock() provider_configurations.get_models.return_value = [ Mock(model="gpt-4", provider=Mock(provider="openai")), @@ -539,16 +613,16 @@ def test_get_first_provider_first_model_returns_first_model_and_provider(mocker: assert result == ("openai", "gpt-4") -def test_update_default_model_record_raises_for_unknown_provider(mocker: MockerFixture): - manager = _build_provider_manager(mocker) +def test_update_default_model_record_raises_for_unknown_provider(): + manager = _build_provider_manager() with patch.object(manager, "get_configurations", return_value={}): with pytest.raises(ValueError, match="Provider openai does not exist."): manager.update_default_model_record("tenant-id", ModelType.LLM, "openai", "gpt-4") -def test_update_default_model_record_raises_for_unknown_model(mocker: MockerFixture): - manager = _build_provider_manager(mocker) +def test_update_default_model_record_raises_for_unknown_model(): + manager = _build_provider_manager() provider_configurations = MagicMock() provider_configurations.__contains__.return_value = True provider_configurations.get_models.return_value = [Mock(model="gpt-4")] @@ -560,8 +634,8 @@ def test_update_default_model_record_raises_for_unknown_model(mocker: MockerFixt provider_configurations.get_models.assert_called_once_with(model_type=ModelType.LLM, only_active=True) -def test_update_default_model_record_updates_existing_record(mocker: MockerFixture): - manager = _build_provider_manager(mocker) +def test_update_default_model_record_updates_existing_record(provider_db: Session): + manager = _build_provider_manager() provider_configurations = MagicMock() provider_configurations.__contains__.return_value = True provider_configurations.get_models.return_value = [Mock(model="gpt-3.5-turbo")] @@ -571,62 +645,87 @@ def test_update_default_model_record_updates_existing_record(mocker: MockerFixtu model_name="claude-3-sonnet", model_type=ModelType.LLM, ) - mock_session = Mock() - mock_session.scalar.return_value = existing_default_model + other_tenant_default = TenantDefaultModel( + tenant_id="other-tenant", + provider_name="cohere", + model_name="command-r", + model_type=ModelType.LLM, + ) + provider_db.add_all([existing_default_model, other_tenant_default]) + provider_db.commit() - with ( - patch.object(manager, "get_configurations", return_value=provider_configurations), - patch("core.provider_manager.db.session", mock_session), - ): + with patch.object(manager, "get_configurations", return_value=provider_configurations): result = manager.update_default_model_record("tenant-id", ModelType.LLM, "openai", "gpt-3.5-turbo") assert result is existing_default_model assert existing_default_model.provider_name == "openai" assert existing_default_model.model_name == "gpt-3.5-turbo" - mock_session.commit.assert_called_once() - mock_session.add.assert_not_called() + provider_db.expire_all() + persisted_default = provider_db.get(TenantDefaultModel, existing_default_model.id) + assert persisted_default is not None + assert persisted_default.provider_name == "openai" + persisted_other_default = provider_db.get(TenantDefaultModel, other_tenant_default.id) + assert persisted_other_default is not None + assert persisted_other_default.provider_name == "cohere" + assert persisted_other_default.model_name == "command-r" -def test_update_default_model_record_creates_record_with_origin_model_type(mocker: MockerFixture): - manager = _build_provider_manager(mocker) +def test_update_default_model_record_creates_record_with_origin_model_type(provider_db: Session): + manager = _build_provider_manager() provider_configurations = MagicMock() provider_configurations.__contains__.return_value = True provider_configurations.get_models.return_value = [Mock(model="gpt-4")] - mock_session = Mock() - mock_session.scalar.return_value = None + provider_db.add( + TenantDefaultModel( + tenant_id="other-tenant", + provider_name="anthropic", + model_name="claude", + model_type=ModelType.LLM, + ) + ) + provider_db.commit() - with ( - patch.object(manager, "get_configurations", return_value=provider_configurations), - patch("core.provider_manager.db.session", mock_session), - ): + with patch.object(manager, "get_configurations", return_value=provider_configurations): result = manager.update_default_model_record("tenant-id", ModelType.LLM, "openai", "gpt-4") - mock_session.add.assert_called_once() - created_default_model = mock_session.add.call_args.args[0] - assert result is created_default_model - assert created_default_model.tenant_id == "tenant-id" - assert created_default_model.provider_name == "openai" - assert created_default_model.model_name == "gpt-4" - assert created_default_model.model_type == ModelType.LLM - mock_session.commit.assert_called_once() + assert result.tenant_id == "tenant-id" + assert result.provider_name == "openai" + assert result.model_name == "gpt-4" + assert result.model_type == ModelType.LLM + persisted_defaults = list(provider_db.scalars(select(TenantDefaultModel)).all()) + assert {record.tenant_id for record in persisted_defaults} == {"tenant-id", "other-tenant"} -def test_get_all_providers_normalizes_provider_names_with_model_provider_id() -> None: - session = Mock() - openai_provider = SimpleNamespace(provider_name="openai") - gemini_provider = SimpleNamespace(provider_name="langgenius/gemini/google") - session.scalars.return_value = [openai_provider, gemini_provider] +def test_get_all_providers_normalizes_provider_names_with_model_provider_id(provider_db: Session) -> None: + openai_provider = Provider( + tenant_id="tenant-id", + provider_name="openai", + provider_type=ProviderType.CUSTOM, + is_valid=True, + ) + gemini_provider = Provider( + tenant_id="tenant-id", + provider_name="langgenius/gemini/google", + provider_type=ProviderType.CUSTOM, + is_valid=True, + ) + other_tenant_provider = Provider( + tenant_id="other-tenant", + provider_name="anthropic", + provider_type=ProviderType.CUSTOM, + is_valid=True, + ) + provider_db.add_all([openai_provider, gemini_provider, other_tenant_provider]) + provider_db.commit() - with ( - patch("core.provider_manager.session_factory.create_session", return_value=_build_session_context(session)), - ): - result = ProviderManager._get_all_providers("tenant-id") + result = ProviderManager._get_all_providers("tenant-id") assert list(result[str(ModelProviderID("openai"))]) == [openai_provider] assert list(result[str(ModelProviderID("langgenius/gemini/google"))]) == [gemini_provider] + assert str(ModelProviderID("anthropic")) not in result -def test_get_all_providers_attaches_active_credentials() -> None: +def test_get_all_providers_attaches_active_credentials(provider_db: Session) -> None: provider = Provider( tenant_id="tenant-id", provider_name="openai", @@ -634,7 +733,6 @@ def test_get_all_providers_attaches_active_credentials() -> None: is_valid=True, credential_id="credential-id", ) - provider.id = "provider-id" credential = ProviderCredential( tenant_id="tenant-id", provider_name="openai", @@ -642,20 +740,25 @@ def test_get_all_providers_attaches_active_credentials() -> None: encrypted_config='{"api_key": "secret"}', ) credential.id = "credential-id" - session = Mock() - session.scalars.side_effect = [ - _FakeScalarResult([provider]), - _FakeScalarResult([credential]), - ] + provider_db.add_all( + [ + provider, + credential, + Provider( + tenant_id="other-tenant", + provider_name="anthropic", + provider_type=ProviderType.CUSTOM, + is_valid=True, + ), + ] + ) + provider_db.commit() - with ( - patch("core.provider_manager.session_factory.create_session", return_value=_build_session_context(session)), - ): - result = ProviderManager._get_all_providers("tenant-id") + result = ProviderManager._get_all_providers("tenant-id") - assert session.scalars.call_count == 2 assert result[str(ModelProviderID("openai"))][0].credential_name == "primary" assert result[str(ModelProviderID("openai"))][0].encrypted_config == '{"api_key": "secret"}' + assert str(ModelProviderID("anthropic")) not in result def test_invalidate_configurations_cache_bumps_selected_source_version() -> None: @@ -676,57 +779,93 @@ def test_invalidate_configurations_cache_bumps_selected_source_version() -> None assert "provider_configurations:tenant:tenant-id:source:provider_models:version" not in fake_redis.store -def test_provider_model_credentials_cache_returns_cache_entries() -> None: +def test_provider_model_credentials_cache_returns_cache_entries(provider_db: Session, sqlite_engine: Engine) -> None: fake_redis = _FakeRedis() - credential_record = SimpleNamespace( - id="credential-id", + credential_record = ProviderModelCredential( + tenant_id="tenant-id", provider_name="openai", model_name="gpt-4", model_type=ModelType.LLM, credential_name="primary", + encrypted_config='{"api_key": "secret"}', ) - session = Mock() - session.scalars.return_value = [credential_record] + credential_record.id = "credential-id" + provider_db.add_all( + [ + credential_record, + ProviderModelCredential( + tenant_id="other-tenant", + provider_name="anthropic", + model_name="claude", + model_type=ModelType.LLM, + credential_name="other", + encrypted_config='{"api_key": "other"}', + ), + ] + ) + provider_db.commit() - with ( - patch("core.provider_manager.redis_client", fake_redis), - patch("core.provider_manager.session_factory.create_session", return_value=_build_session_context(session)), - ): - first = ProviderManager._get_all_provider_model_credentials("tenant-id") - second = ProviderManager._get_all_provider_model_credentials("tenant-id") + credential_query_count = 0 + + def count_credential_queries(_conn, _cursor, statement, _parameters, _context, _executemany) -> None: + nonlocal credential_query_count + if statement.lstrip().upper().startswith("SELECT") and "provider_model_credentials" in statement: + credential_query_count += 1 + + # Count real SQL statements now that this test uses SQLite instead of a mocked session. + event.listen(sqlite_engine, "before_cursor_execute", count_credential_queries) + try: + with patch("core.provider_manager.redis_client", fake_redis): + first = ProviderManager._get_all_provider_model_credentials("tenant-id") + second = ProviderManager._get_all_provider_model_credentials("tenant-id") + finally: + event.remove(sqlite_engine, "before_cursor_execute", count_credential_queries) - assert session.scalars.call_count == 1 version_key = "provider_configurations:tenant:tenant-id:source:provider_model_credentials:version" + assert credential_query_count == 1 assert fake_redis.expirations[version_key] == 360 assert first["openai"][0] is not credential_record assert second["openai"][0].credential_name == "primary" assert second["openai"][0].model_type == ModelType.LLM + assert "anthropic" not in second -def test_provider_configuration_cache_skips_write_when_version_changes_during_load() -> None: +def test_provider_configuration_cache_skips_write_when_version_changes_during_load( + provider_db: Session, sqlite_engine: Engine +) -> None: fake_redis = _FakeRedis() version_key = "provider_configurations:tenant:tenant-id:source:provider_model_credentials:version" - credential_record = SimpleNamespace( - id="credential-id", + credential_record = ProviderModelCredential( + tenant_id="tenant-id", provider_name="openai", model_name="gpt-4", model_type=ModelType.LLM, credential_name="primary", + encrypted_config='{"api_key": "secret"}', ) - session = Mock() + credential_record.id = "credential-id" + provider_db.add(credential_record) + provider_db.commit() + version_bumped = False - def load_records(_stmt): - fake_redis.incr(version_key) - return [credential_record] + def bump_version_on_credential_query(_conn, _cursor, statement, _parameters, _context, _executemany) -> None: + nonlocal version_bumped + if ( + not version_bumped + and statement.lstrip().upper().startswith("SELECT") + and "provider_model_credentials" in statement + ): + version_bumped = True + fake_redis.incr(version_key) - session.scalars.side_effect = load_records - - with ( - patch("core.provider_manager.redis_client", fake_redis), - patch("core.provider_manager.session_factory.create_session", return_value=_build_session_context(session)), - ): - result = ProviderManager._get_all_provider_model_credentials("tenant-id") + event.listen(sqlite_engine, "before_cursor_execute", bump_version_on_credential_query) + try: + with patch("core.provider_manager.redis_client", fake_redis): + result = ProviderManager._get_all_provider_model_credentials("tenant-id") + finally: + event.remove(sqlite_engine, "before_cursor_execute", bump_version_on_credential_query) + assert version_bumped is True assert fake_redis.store[version_key] == "1" assert "provider_configurations:tenant:tenant-id:source:provider_model_credentials:v:0" not in fake_redis.store assert "provider_configurations:tenant:tenant-id:source:provider_model_credentials:v:1" not in fake_redis.store @@ -742,19 +881,21 @@ def test_provider_configuration_cache_skips_write_when_version_changes_during_lo "_get_all_provider_credentials", ], ) -def test_provider_grouping_helpers_group_records_by_provider_name(method_name: str) -> None: - def build_record(provider_name: str, index: int): +def test_provider_grouping_helpers_group_records_by_provider_name(method_name: str, provider_db: Session) -> None: + def build_record(provider_name: str, index: int, *, tenant_id: str = "tenant-id"): match method_name: case "_get_all_provider_models": - return SimpleNamespace( - id=f"model-{index}", + record = ProviderModel( + tenant_id=tenant_id, provider_name=provider_name, model_name=f"model-{index}", model_type=ModelType.LLM, credential_id=None, + is_valid=True, ) case "_get_all_provider_model_settings": - return SimpleNamespace( + record = ProviderModelSetting( + tenant_id=tenant_id, provider_name=provider_name, model_name=f"model-{index}", model_type=ModelType.LLM, @@ -762,69 +903,95 @@ def test_provider_grouping_helpers_group_records_by_provider_name(method_name: s load_balancing_enabled=False, ) case "_get_all_provider_model_credentials": - return SimpleNamespace( - id=f"model-credential-{index}", + record = ProviderModelCredential( + tenant_id=tenant_id, provider_name=provider_name, model_name=f"model-{index}", model_type=ModelType.LLM, credential_name=f"credential-{index}", + encrypted_config='{"api_key": "secret"}', ) case "_get_all_provider_credentials": - return SimpleNamespace( - id=f"credential-{index}", + record = ProviderCredential( + tenant_id=tenant_id, provider_name=provider_name, credential_name=f"credential-{index}", + encrypted_config='{"api_key": "secret"}', ) case _: raise AssertionError(f"Unexpected method: {method_name}") + record.id = f"record-{index}" + return record - session = Mock() openai_primary = build_record("openai", 1) openai_secondary = build_record("openai", 2) anthropic_record = build_record("anthropic", 3) - session.scalars.return_value = [openai_primary, openai_secondary, anthropic_record] + other_tenant_record = build_record("other-provider", 4, tenant_id="other-tenant") + provider_db.add_all([openai_primary, openai_secondary, anthropic_record, other_tenant_record]) + provider_db.commit() - with ( - patch("core.provider_manager.session_factory.create_session", return_value=_build_session_context(session)), - ): + with patch("core.provider_manager.redis_client", _FakeRedis()): result = getattr(ProviderManager, method_name)("tenant-id") assert [record.provider_name for record in result["openai"]] == ["openai", "openai"] assert [record.provider_name for record in result["anthropic"]] == ["anthropic"] + assert "other-provider" not in result -def test_get_all_preferred_model_providers_returns_mapping_by_provider_name() -> None: - session = Mock() - openai_preference = SimpleNamespace(provider_name="openai", preferred_provider_type=ProviderType.SYSTEM) - anthropic_preference = SimpleNamespace(provider_name="anthropic", preferred_provider_type=ProviderType.CUSTOM) - session.scalars.return_value = [openai_preference, anthropic_preference] +def test_get_all_preferred_model_providers_returns_mapping_by_provider_name(provider_db: Session) -> None: + openai_preference = TenantPreferredModelProvider( + tenant_id="tenant-id", provider_name="openai", preferred_provider_type=ProviderType.SYSTEM + ) + anthropic_preference = TenantPreferredModelProvider( + tenant_id="tenant-id", provider_name="anthropic", preferred_provider_type=ProviderType.CUSTOM + ) + provider_db.add_all( + [ + openai_preference, + anthropic_preference, + TenantPreferredModelProvider( + tenant_id="other-tenant", + provider_name="other-provider", + preferred_provider_type=ProviderType.SYSTEM, + ), + ] + ) + provider_db.commit() - with ( - patch("core.provider_manager.session_factory.create_session", return_value=_build_session_context(session)), - ): + with patch("core.provider_manager.redis_client", _FakeRedis()): result = ProviderManager._get_all_preferred_model_providers("tenant-id") assert result["openai"].preferred_provider_type == ProviderType.SYSTEM assert result["anthropic"].preferred_provider_type == ProviderType.CUSTOM + assert "other-provider" not in result -def test_get_all_provider_load_balancing_configs_returns_empty_when_cached_flag_is_disabled() -> None: - with ( - patch("core.provider_manager.redis_client.get", return_value=b"False"), - patch("core.provider_manager.FeatureService.get_features") as mock_get_features, - patch("core.provider_manager.session_factory.create_session") as mock_create_session, - ): - result = ProviderManager._get_all_provider_load_balancing_configs("tenant-id") +@pytest.mark.usefixtures("provider_db") +def test_get_all_provider_load_balancing_configs_returns_empty_when_cached_flag_is_disabled( + sqlite_engine: Engine, +) -> None: + statements: list[str] = [] + + def record_statement(_conn, _cursor, statement, _parameters, _context, _executemany) -> None: + statements.append(statement) + + event.listen(sqlite_engine, "before_cursor_execute", record_statement) + try: + with ( + patch("core.provider_manager.redis_client.get", return_value=b"False"), + patch("core.provider_manager.FeatureService.get_features") as mock_get_features, + ): + result = ProviderManager._get_all_provider_load_balancing_configs("tenant-id") + finally: + event.remove(sqlite_engine, "before_cursor_execute", record_statement) assert result == {} mock_get_features.assert_not_called() - mock_create_session.assert_not_called() + assert statements == [] -def test_get_all_provider_load_balancing_configs_populates_cache_and_groups_configs() -> None: - session = Mock() - openai_config = SimpleNamespace( - id="lb-1", +def test_get_all_provider_load_balancing_configs_populates_cache_and_groups_configs(provider_db: Session) -> None: + openai_config = LoadBalancingModelConfig( tenant_id="tenant-id", provider_name="openai", model_name="gpt-4", @@ -835,8 +1002,8 @@ def test_get_all_provider_load_balancing_configs_populates_cache_and_groups_conf credential_source_type=None, enabled=True, ) - anthropic_config = SimpleNamespace( - id="lb-2", + openai_config.id = "lb-1" + anthropic_config = LoadBalancingModelConfig( tenant_id="tenant-id", provider_name="anthropic", model_name="claude", @@ -847,7 +1014,20 @@ def test_get_all_provider_load_balancing_configs_populates_cache_and_groups_conf credential_source_type=None, enabled=True, ) - session.scalars.return_value = [openai_config, anthropic_config] + anthropic_config.id = "lb-2" + other_tenant_config = LoadBalancingModelConfig( + tenant_id="other-tenant", + provider_name="other-provider", + model_name="other-model", + model_type=ModelType.LLM, + name="primary", + encrypted_config=None, + credential_id=None, + credential_source_type=None, + enabled=True, + ) + provider_db.add_all([openai_config, anthropic_config, other_tenant_config]) + provider_db.commit() with ( patch("core.provider_manager.redis_client.get", return_value=None), @@ -856,10 +1036,10 @@ def test_get_all_provider_load_balancing_configs_populates_cache_and_groups_conf "core.provider_manager.FeatureService.get_features", return_value=SimpleNamespace(model_load_balancing_enabled=True), ), - patch("core.provider_manager.session_factory.create_session", return_value=_build_session_context(session)), ): result = ProviderManager._get_all_provider_load_balancing_configs("tenant-id") mock_setex.assert_any_call("tenant:tenant-id:model_load_balancing_enabled", 120, "True") assert [record.provider_name for record in result["openai"]] == ["openai"] assert [record.provider_name for record in result["anthropic"]] == ["anthropic"] + assert "other-provider" not in result