diff --git a/api/tests/unit_tests/core/entities/test_entities_provider_configuration.py b/api/tests/unit_tests/core/entities/test_entities_provider_configuration.py index b5489c88a03..40072487062 100644 --- a/api/tests/unit_tests/core/entities/test_entities_provider_configuration.py +++ b/api/tests/unit_tests/core/entities/test_entities_provider_configuration.py @@ -5,7 +5,6 @@ import logging from collections.abc import Iterator from contextlib import contextmanager from types import SimpleNamespace -from typing import Any from unittest.mock import Mock, PropertyMock, patch import pytest @@ -55,8 +54,6 @@ from models.provider import ( ) from models.provider_ids import ModelProviderID -_UNSET = object() - def _build_provider_configuration(*, provider_name: str = "openai") -> ProviderConfiguration: provider_entity = ProviderEntity( @@ -102,37 +99,6 @@ def _build_ai_model(name: str, *, model_type: ModelType = ModelType.LLM) -> AIMo ) -def _exec_result( - *, - scalar_one_or_none: Any = _UNSET, - scalar: Any = _UNSET, - scalars_all: Any = _UNSET, - scalars_first: Any = _UNSET, -) -> Mock: - result = Mock() - if scalar_one_or_none is not _UNSET: - result.scalar_one_or_none.return_value = scalar_one_or_none - if scalar is not _UNSET: - result.scalar.return_value = scalar - if scalars_all is not _UNSET or scalars_first is not _UNSET: - scalars = Mock() - if scalars_all is not _UNSET: - scalars.all.return_value = scalars_all - if scalars_first is not _UNSET: - scalars.first.return_value = scalars_first - result.scalars.return_value = scalars - return result - - -@contextmanager -def _patched_session(session: Mock): - with patch("core.entities.provider_configuration.db") as mock_db: - mock_db.engine = Mock() - with patch("core.entities.provider_configuration.Session") as mock_session_cls: - mock_session_cls.return_value.__enter__.return_value = session - yield mock_session_cls - - def _build_secret_provider_schema() -> ProviderCredentialSchema: return ProviderCredentialSchema( credential_form_schemas=[ @@ -300,158 +266,6 @@ def test_get_provider_names_supports_legacy_and_full_plugin_id() -> None: assert provider_names == ["langgenius/openai/openai", "openai"] -def test_generate_next_api_key_name_uses_highest_numeric_suffix() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.execute.return_value.scalars.return_value.all.return_value = [ - SimpleNamespace(credential_name="API KEY 9"), - SimpleNamespace(credential_name="legacy"), - SimpleNamespace(credential_name=" API KEY 2 "), - ] - - name = configuration._generate_next_api_key_name(session=session, query_factory=lambda: Mock()) - assert name == "API KEY 10" - - -def test_generate_next_api_key_name_falls_back_to_default_on_error() -> None: - configuration = _build_provider_configuration() - session = Mock() - - def _raise_query_error(): - raise RuntimeError("boom") - - name = configuration._generate_next_api_key_name(session=session, query_factory=_raise_query_error) - assert name == "API KEY 1" - - -def test_generate_provider_and_custom_model_names_delegate_to_shared_generator() -> None: - configuration = _build_provider_configuration() - - with patch.object(configuration, "_generate_next_api_key_name", return_value="API KEY 7") as mock_generator: - provider_name = configuration._generate_provider_credential_name(session=Mock()) - custom_model_name = configuration._generate_custom_model_credential_name( - model="gpt-4o", - model_type=ModelType.LLM, - session=Mock(), - ) - - assert provider_name == "API KEY 7" - assert custom_model_name == "API KEY 7" - assert mock_generator.call_count == 2 - - -def test_get_provider_credential_uses_specific_lookup_when_id_provided() -> None: - configuration = _build_provider_configuration() - - with patch.object(configuration, "_get_specific_provider_credential", return_value={"api_key": "***"}) as mock_get: - credential = configuration.get_provider_credential("credential-1") - - assert credential == {"api_key": "***"} - mock_get.assert_called_once_with("credential-1") - - -def test_validate_provider_credentials_handles_hidden_secret_value() -> None: - configuration = _build_provider_configuration() - configuration.provider.provider_credential_schema = ProviderCredentialSchema( - credential_form_schemas=[ - CredentialFormSchema( - variable="openai_api_key", - label=I18nObject(en_US="API Key"), - type=FormType.SECRET_INPUT, - ) - ] - ) - mock_session = Mock() - mock_session.execute.return_value.scalar_one_or_none.return_value = SimpleNamespace( - encrypted_config="encrypted-old-key" - ) - mock_factory = Mock() - mock_factory.provider_credentials_validate.return_value = {"openai_api_key": "restored-key", "region": "us"} - - with _patched_session(mock_session): - with patch( - "core.entities.provider_configuration.create_plugin_model_assembly", - return_value=SimpleNamespace(model_runtime=Mock(), model_provider_factory=mock_factory), - ): - with patch("core.entities.provider_configuration.encrypter.decrypt_token", return_value="restored-key"): - with patch( - "core.entities.provider_configuration.encrypter.encrypt_token", - side_effect=lambda tenant_id, value: f"enc::{value}", - ): - validated = configuration.validate_provider_credentials( - credentials={"openai_api_key": HIDDEN_VALUE, "region": "us"}, - credential_id="credential-1", - ) - - assert validated["openai_api_key"] == "enc::restored-key" - assert validated["region"] == "us" - mock_factory.provider_credentials_validate.assert_called_once_with( - provider="openai", - credentials={"openai_api_key": "restored-key", "region": "us"}, - ) - - -def test_validate_provider_credentials_without_credential_id() -> None: - configuration = _build_provider_configuration() - mock_factory = Mock() - mock_factory.provider_credentials_validate.return_value = {"region": "us"} - - with patch( - "core.entities.provider_configuration.create_plugin_model_assembly", - return_value=SimpleNamespace(model_runtime=Mock(), model_provider_factory=mock_factory), - ): - validated = configuration.validate_provider_credentials(credentials={"region": "us"}) - - assert validated == {"region": "us"} - - -def test_switch_preferred_provider_type_returns_early_when_no_change_or_unsupported() -> None: - configuration = _build_provider_configuration() - - with patch("core.entities.provider_configuration.Session") as mock_session_cls: - changed = configuration.switch_preferred_provider_type(ProviderType.SYSTEM) - assert changed is False - mock_session_cls.assert_not_called() - - configuration.preferred_provider_type = ProviderType.CUSTOM - configuration.system_configuration.enabled = False - with patch("core.entities.provider_configuration.Session") as mock_session_cls: - changed = configuration.switch_preferred_provider_type(ProviderType.SYSTEM) - assert changed is False - mock_session_cls.assert_not_called() - - -def test_switch_preferred_provider_type_updates_existing_record_with_session() -> None: - configuration = _build_provider_configuration() - configuration.preferred_provider_type = ProviderType.CUSTOM - session = Mock() - existing_record = SimpleNamespace(preferred_provider_type="custom") - session.execute.return_value.scalars.return_value.first.return_value = existing_record - - with patch.object(ProviderConfiguration, "_invalidate_provider_configuration_cache") as mock_invalidate: - changed = configuration.switch_preferred_provider_type(ProviderType.SYSTEM, session=session) - - assert changed is True - assert existing_record.preferred_provider_type == ProviderType.SYSTEM - session.commit.assert_called_once() - mock_invalidate.assert_not_called() - - -def test_switch_preferred_provider_type_creates_record_when_missing() -> None: - configuration = _build_provider_configuration() - configuration.preferred_provider_type = ProviderType.SYSTEM - session = Mock() - session.execute.return_value.scalars.return_value.first.return_value = None - - with patch.object(ProviderConfiguration, "_invalidate_provider_configuration_cache") as mock_invalidate: - changed = configuration.switch_preferred_provider_type(ProviderType.CUSTOM, session=session) - - assert changed is True - assert session.add.call_count == 1 - session.commit.assert_called_once() - mock_invalidate.assert_not_called() - - def test_get_model_type_instance_and_schema_delegate_to_factory() -> None: configuration = _build_provider_configuration() mock_model_type_instance = Mock() @@ -851,1269 +665,6 @@ def test_get_system_configuration_status_falsey_quota_returns_unsupported() -> N assert configuration.get_system_configuration_status() == SystemConfigurationStatus.UNSUPPORTED -def test_get_provider_credential_default_uses_custom_provider_credentials() -> None: - configuration = _build_provider_configuration() - configuration.custom_configuration.provider = CustomProviderConfiguration(credentials={"api_key": "provider-key"}) - obfuscated = configuration.get_provider_credential() - assert obfuscated == {"api_key": "provider-key"} - - -def test_custom_configuration_availability_and_provider_record_helpers() -> None: - configuration = _build_provider_configuration() - assert not configuration.is_custom_configuration_available() - - configuration.custom_configuration.provider = CustomProviderConfiguration( - credentials={"api_key": "provider-key"}, - available_credentials=[CredentialConfiguration(credential_id="cred-1", credential_name="Main")], - ) - assert configuration.is_custom_configuration_available() - - configuration.custom_configuration.provider = None - configuration.custom_configuration.models = [ - CustomModelConfiguration(model="gpt-4o", model_type=ModelType.LLM, credentials={"api_key": "model-key"}) - ] - assert configuration.is_custom_configuration_available() - - session = Mock() - provider_record = SimpleNamespace(id="provider-1") - session.execute.return_value.scalar_one_or_none.return_value = provider_record - assert configuration._get_provider_record(session) is provider_record - - session.execute.return_value.scalar_one_or_none.return_value = None - assert configuration._get_provider_record(session) is None - - -def test_check_provider_credential_name_exists_and_model_setting_lookup() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = "existing-id" - assert configuration._check_provider_credential_name_exists("Main", session) - - session.execute.return_value.scalar_one_or_none.return_value = None - assert not configuration._check_provider_credential_name_exists("Main", session, exclude_id="cred-2") - - setting = SimpleNamespace(id="setting-1") - session.execute.return_value.scalars.return_value.first.return_value = setting - assert configuration._get_provider_model_setting(ModelType.LLM, "gpt-4o", session) is setting - - -def test_validate_provider_credentials_handles_invalid_original_json() -> None: - configuration = _build_provider_configuration() - configuration.provider.provider_credential_schema = _build_secret_provider_schema() - mock_session = Mock() - mock_session.execute.return_value.scalar_one_or_none.return_value = SimpleNamespace( - encrypted_config="{invalid-json" - ) - mock_factory = Mock() - mock_factory.provider_credentials_validate.return_value = {"openai_api_key": "new-key"} - - with _patched_session(mock_session): - with patch( - "core.entities.provider_configuration.create_plugin_model_assembly", - return_value=SimpleNamespace(model_runtime=Mock(), model_provider_factory=mock_factory), - ): - with patch("core.entities.provider_configuration.encrypter.encrypt_token", return_value="enc-key"): - validated = configuration.validate_provider_credentials( - credentials={"openai_api_key": HIDDEN_VALUE}, - credential_id="cred-1", - ) - - assert validated == {"openai_api_key": "enc-key"} - - -def test_generate_next_api_key_name_returns_default_when_no_records() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.execute.return_value.scalars.return_value.all.return_value = [] - - name = configuration._generate_next_api_key_name(session=session, query_factory=lambda: Mock()) - assert name == "API KEY 1" - - -def test_create_provider_credential_creates_provider_record_when_missing() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.flush.side_effect = lambda: None - - with _patched_session(session): - with patch.object(ProviderConfiguration, "validate_provider_credentials", return_value={"api_key": "enc"}): - with patch.object(ProviderConfiguration, "_get_provider_record", return_value=None): - with patch.object( - ProviderConfiguration, - "_generate_provider_credential_name", - return_value="API KEY 2", - ): - with patch.object(ProviderConfiguration, "switch_preferred_provider_type") as mock_switch: - with patch("core.entities.provider_configuration.ProviderCredentialsCache") as mock_cache: - configuration.create_provider_credential({"api_key": "raw"}, None) - - assert session.add.call_count == 2 - session.commit.assert_called_once() - mock_cache.return_value.delete.assert_called_once() - mock_switch.assert_called_once_with(provider_type=ProviderType.CUSTOM, session=session) - - -def test_create_provider_credential_marks_existing_provider_as_valid() -> None: - configuration = _build_provider_configuration() - session = Mock() - provider_record = SimpleNamespace(id="provider-1", is_valid=False, credential_id="existing-cred") - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_check_provider_credential_name_exists", return_value=False): - with patch.object(ProviderConfiguration, "validate_provider_credentials", return_value={"api_key": "enc"}): - with patch.object(ProviderConfiguration, "_get_provider_record", return_value=provider_record): - configuration.create_provider_credential({"api_key": "raw"}, "Main") - - assert provider_record.is_valid is True - assert provider_record.credential_id == "existing-cred" - session.commit.assert_called_once() - - -def test_create_provider_credential_auto_activates_when_no_active_credential() -> None: - configuration = _build_provider_configuration() - session = Mock() - provider_record = SimpleNamespace(id="provider-1", is_valid=False, credential_id=None, updated_at=None) - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_check_provider_credential_name_exists", return_value=False): - with patch.object(ProviderConfiguration, "validate_provider_credentials", return_value={"api_key": "enc"}): - with patch.object(ProviderConfiguration, "_get_provider_record", return_value=provider_record): - with patch("core.entities.provider_configuration.ProviderCredentialsCache"): - with patch.object(ProviderConfiguration, "switch_preferred_provider_type"): - configuration.create_provider_credential({"api_key": "raw"}, "Main") - - assert provider_record.is_valid is True - assert provider_record.credential_id is not None - session.commit.assert_called_once() - - -def test_create_provider_credential_raises_when_duplicate_name_exists() -> None: - configuration = _build_provider_configuration() - session = Mock() - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_check_provider_credential_name_exists", return_value=True): - with pytest.raises(ValueError, match="already exists"): - configuration.create_provider_credential({"api_key": "raw"}, "Main") - - -def test_update_provider_credential_success_updates_and_invalidates_cache() -> None: - configuration = _build_provider_configuration() - session = Mock() - credential_record = SimpleNamespace(id="cred-1", encrypted_config="{}", credential_name="Old", updated_at=None) - provider_record = SimpleNamespace(id="provider-1", credential_id="cred-1") - session.execute.return_value.scalar_one_or_none.return_value = credential_record - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_check_provider_credential_name_exists", return_value=False): - with patch.object(ProviderConfiguration, "validate_provider_credentials", return_value={"api_key": "enc"}): - with patch.object(ProviderConfiguration, "_get_provider_record", return_value=provider_record): - with patch.object( - ProviderConfiguration, - "_update_load_balancing_configs_with_credential", - ) as mock_lb: - with patch("core.entities.provider_configuration.ProviderCredentialsCache") as mock_cache: - configuration.update_provider_credential( - credentials={"api_key": "raw"}, - credential_id="cred-1", - credential_name="New Name", - ) - - assert credential_record.credential_name == "New Name" - session.commit.assert_called_once() - mock_cache.return_value.delete.assert_called_once() - mock_lb.assert_called_once() - - -def test_update_provider_credential_raises_when_record_not_found() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = None - - with _patched_session(session): - with patch.object(ProviderConfiguration, "validate_provider_credentials", return_value={"api_key": "enc"}): - with patch.object(ProviderConfiguration, "_get_provider_record", return_value=None): - with pytest.raises(ValueError, match="Credential record not found"): - configuration.update_provider_credential({"api_key": "raw"}, "cred-1", None) - - -def test_update_load_balancing_configs_updates_all_matching_configs() -> None: - configuration = _build_provider_configuration() - session = Mock() - lb_config = SimpleNamespace(id="lb-1", encrypted_config="old", name="old", updated_at=None) - session.execute.return_value.scalars.return_value.all.return_value = [lb_config] - credential_record = SimpleNamespace(encrypted_config='{"api_key":"enc"}', credential_name="API KEY 3") - - with patch("core.entities.provider_configuration.ProviderCredentialsCache") as mock_cache: - changed = configuration._update_load_balancing_configs_with_credential( - credential_id="cred-1", - credential_record=credential_record, - credential_source=CredentialSourceType.PROVIDER, - session=session, - ) - - assert changed is True - assert lb_config.encrypted_config == '{"api_key":"enc"}' - assert lb_config.name == "API KEY 3" - mock_cache.return_value.delete.assert_called_once() - session.commit.assert_called_once() - - -def test_update_load_balancing_configs_returns_when_no_matching_configs() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.execute.return_value.scalars.return_value.all.return_value = [] - - changed = configuration._update_load_balancing_configs_with_credential( - credential_id="cred-1", - credential_record=SimpleNamespace(encrypted_config="{}", credential_name="Main"), - credential_source=CredentialSourceType.PROVIDER, - session=session, - ) - - assert changed is False - session.commit.assert_not_called() - - -def test_delete_provider_credential_removes_provider_record_when_last_credential() -> None: - configuration = _build_provider_configuration() - session = Mock() - credential_record = SimpleNamespace(id="cred-1") - provider_record = SimpleNamespace(id="provider-1", credential_id="cred-1", updated_at=None) - session.execute.side_effect = [ - _exec_result(scalar_one_or_none=credential_record), - _exec_result(scalars_all=[]), - _exec_result(scalar=1), - ] - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_record", return_value=provider_record): - with patch.object(ProviderConfiguration, "switch_preferred_provider_type") as mock_switch: - with patch("core.entities.provider_configuration.ProviderCredentialsCache") as mock_cache: - configuration.delete_provider_credential("cred-1") - - assert any(call.args and call.args[0] is provider_record for call in session.delete.call_args_list) - mock_cache.return_value.delete.assert_called_once() - mock_switch.assert_called_once_with(provider_type=ProviderType.SYSTEM, session=session) - - -def test_delete_provider_credential_raises_when_not_found() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = None - - with _patched_session(session): - with pytest.raises(ValueError, match="Credential record not found"): - configuration.delete_provider_credential("cred-1") - - -def test_delete_provider_credential_unsets_active_credential_when_more_available() -> None: - configuration = _build_provider_configuration() - session = Mock() - credential_record = SimpleNamespace(id="cred-1") - lb_config = SimpleNamespace(id="lb-1") - provider_record = SimpleNamespace(id="provider-1", credential_id="cred-1", updated_at=None) - session.execute.side_effect = [ - _exec_result(scalar_one_or_none=credential_record), - _exec_result(scalars_all=[lb_config]), - _exec_result(scalar=2), - ] - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_record", return_value=provider_record): - with patch.object(ProviderConfiguration, "switch_preferred_provider_type") as mock_switch: - with patch("core.entities.provider_configuration.ProviderCredentialsCache") as mock_cache: - configuration.delete_provider_credential("cred-1") - - assert provider_record.credential_id is None - assert mock_cache.return_value.delete.call_count == 2 - mock_switch.assert_called_once_with(provider_type=ProviderType.SYSTEM, session=session) - - -def test_switch_active_provider_credential_success_and_failures() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = None - with _patched_session(session): - with pytest.raises(ValueError, match="Credential record not found"): - configuration.switch_active_provider_credential("cred-1") - - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = SimpleNamespace(id="cred-1") - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_record", return_value=None): - with pytest.raises(ValueError, match="Provider record not found"): - configuration.switch_active_provider_credential("cred-1") - - session = Mock() - credential_record = SimpleNamespace(id="cred-1") - provider_record = SimpleNamespace(id="provider-1", credential_id=None, updated_at=None) - session.execute.return_value.scalar_one_or_none.return_value = credential_record - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_record", return_value=provider_record): - with patch.object(ProviderConfiguration, "switch_preferred_provider_type") as mock_switch: - with patch("core.entities.provider_configuration.ProviderCredentialsCache") as mock_cache: - configuration.switch_active_provider_credential("cred-1") - - assert provider_record.credential_id == "cred-1" - mock_cache.return_value.delete.assert_called_once() - mock_switch.assert_called_once_with(ProviderType.CUSTOM, session=session) - - -def test_get_custom_model_record_supports_plugin_id_alias() -> None: - configuration = _build_provider_configuration(provider_name="langgenius/openai/openai") - session = Mock() - custom_model_record = SimpleNamespace(id="model-1") - session.execute.return_value.scalar_one_or_none.return_value = custom_model_record - - result = configuration._get_custom_model_record(ModelType.LLM, "gpt-4o", session) - assert result is custom_model_record - - -def test_model_type_db_values_includes_pre_1_15_aliases() -> None: - from core.entities.provider_configuration import _model_type_db_values - - assert _model_type_db_values(ModelType.LLM) == ("llm", "text-generation") - assert _model_type_db_values(ModelType.TEXT_EMBEDDING) == ("text-embedding", "embeddings") - assert _model_type_db_values(ModelType.RERANK) == ("rerank", "reranking") - assert _model_type_db_values(ModelType.TTS) == ("tts",) - - -def test_get_custom_model_record_uses_legacy_aware_model_type_filter( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Regression for #39559: lookups must include pre-1.15 model_type aliases.""" - import core.entities.provider_configuration as provider_configuration_module - - captured: dict[str, tuple[str, ...]] = {} - original = provider_configuration_module._model_type_db_values - - def _capture(model_type: ModelType) -> tuple[str, ...]: - values = original(model_type) - captured["values"] = values - return values - - monkeypatch.setattr(provider_configuration_module, "_model_type_db_values", _capture) - - configuration = _build_provider_configuration(provider_name="langgenius/ollama/ollama") - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = SimpleNamespace(id="legacy-model") - - result = configuration._get_custom_model_record(ModelType.LLM, "llama3", session) - - assert result.id == "legacy-model" - assert captured["values"] == ("llm", "text-generation") - - -def test_get_specific_custom_model_credential_success_and_not_found() -> None: - configuration = _build_provider_configuration() - configuration.provider.model_credential_schema = _build_secret_model_schema() - session = Mock() - record = SimpleNamespace(id="cred-1", credential_name="Main", encrypted_config='{"openai_api_key":"enc"}') - session.execute.return_value.scalar_one_or_none.return_value = record - - with _patched_session(session): - with patch("core.entities.provider_configuration.encrypter.decrypt_token", return_value="raw"): - with patch.object(ProviderConfiguration, "obfuscated_credentials", return_value={"openai_api_key": "***"}): - response = configuration._get_specific_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") - - assert response["current_credential_id"] == "cred-1" - assert response["credentials"] == {"openai_api_key": "***"} - - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = None - with _patched_session(session): - with pytest.raises(ValueError, match="Credential with id cred-1 not found"): - configuration._get_specific_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") - - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = SimpleNamespace( - id="cred-1", - credential_name="Main", - encrypted_config="{invalid-json", - ) - with _patched_session(session): - invalid_json = configuration._get_specific_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") - assert invalid_json["credentials"] == {} - - -def test_check_custom_model_credential_name_exists_respects_exclusion() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = SimpleNamespace(id="cred-1") - assert configuration._check_custom_model_credential_name_exists( - ModelType.LLM, "gpt-4o", "Main", session, exclude_id="other-id" - ) - - session.execute.return_value.scalar_one_or_none.return_value = None - assert not configuration._check_custom_model_credential_name_exists(ModelType.LLM, "gpt-4o", "Main", session) - - -def test_get_custom_model_credential_uses_specific_id_or_configuration_fallback() -> None: - configuration = _build_provider_configuration() - with patch.object( - ProviderConfiguration, - "_get_specific_custom_model_credential", - return_value={"current_credential_id": "cred-1"}, - ) as mock_specific: - result = configuration.get_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") - assert result == {"current_credential_id": "cred-1"} - mock_specific.assert_called_once() - - configuration.provider.model_credential_schema = _build_secret_model_schema() - configuration.custom_configuration.models = [ - CustomModelConfiguration( - model="gpt-4o", - model_type=ModelType.LLM, - credentials={"openai_api_key": "raw"}, - current_credential_id="cred-1", - current_credential_name="Main", - ) - ] - with patch.object(ProviderConfiguration, "obfuscated_credentials", return_value={"openai_api_key": "***"}): - fallback = configuration.get_custom_model_credential(ModelType.LLM, "gpt-4o", None) - assert fallback == { - "current_credential_id": "cred-1", - "current_credential_name": "Main", - "credentials": {"openai_api_key": "***"}, - } - - configuration.custom_configuration.models = [] - assert configuration.get_custom_model_credential(ModelType.LLM, "gpt-4o", None) is None - - -def test_validate_custom_model_credentials_supports_hidden_reuse_and_sessionless_path() -> None: - configuration = _build_provider_configuration() - configuration.provider.model_credential_schema = _build_secret_model_schema() - mock_session = Mock() - mock_session.execute.return_value.scalar_one_or_none.return_value = SimpleNamespace( - encrypted_config='{"openai_api_key":"enc"}' - ) - mock_factory = Mock() - mock_factory.model_credentials_validate.return_value = {"openai_api_key": "raw"} - - with _patched_session(mock_session): - with patch( - "core.entities.provider_configuration.create_plugin_model_assembly", - return_value=SimpleNamespace(model_runtime=Mock(), model_provider_factory=mock_factory), - ): - with patch("core.entities.provider_configuration.encrypter.decrypt_token", return_value="raw"): - with patch("core.entities.provider_configuration.encrypter.encrypt_token", return_value="enc-new"): - validated = configuration.validate_custom_model_credentials( - model_type=ModelType.LLM, - model="gpt-4o", - credentials={"openai_api_key": HIDDEN_VALUE}, - credential_id="cred-1", - ) - assert validated == {"openai_api_key": "enc-new"} - - mock_factory2 = Mock() - mock_factory2.model_credentials_validate.return_value = {"region": "us"} - with patch( - "core.entities.provider_configuration.create_plugin_model_assembly", - return_value=SimpleNamespace(model_runtime=Mock(), model_provider_factory=mock_factory2), - ): - validated = configuration.validate_custom_model_credentials( - model_type=ModelType.LLM, - model="gpt-4o", - credentials={"region": "us"}, - ) - assert validated == {"region": "us"} - - -def test_create_update_delete_custom_model_credential_flow() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.flush.side_effect = lambda: None - provider_model_record = SimpleNamespace(id="model-1", credential_id="cred-1", updated_at=None) - credential_record = SimpleNamespace(id="cred-1", encrypted_config="{}", credential_name="Old", updated_at=None) - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_generate_custom_model_credential_name", return_value="API KEY 1"): - with patch.object( - ProviderConfiguration, - "validate_custom_model_credentials", - return_value={"openai_api_key": "enc"}, - ): - with patch.object(ProviderConfiguration, "_get_custom_model_record", return_value=None): - with patch("core.entities.provider_configuration.ProviderCredentialsCache") as mock_cache: - configuration.create_custom_model_credential(ModelType.LLM, "gpt-4o", {"k": "v"}, None) - assert session.add.call_count == 2 - assert mock_cache.return_value.delete.call_count == 1 - - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = credential_record - with _patched_session(session): - with patch.object(ProviderConfiguration, "_check_custom_model_credential_name_exists", return_value=False): - with patch.object( - ProviderConfiguration, - "validate_custom_model_credentials", - return_value={"openai_api_key": "enc2"}, - ): - with patch.object( - ProviderConfiguration, - "_get_custom_model_record", - return_value=provider_model_record, - ): - with patch.object( - ProviderConfiguration, - "_update_load_balancing_configs_with_credential", - ) as mock_lb: - with patch("core.entities.provider_configuration.ProviderCredentialsCache") as mock_cache: - configuration.update_custom_model_credential( - model_type=ModelType.LLM, - model="gpt-4o", - credentials={"k": "v"}, - credential_name="New Name", - credential_id="cred-1", - ) - assert credential_record.credential_name == "New Name" - assert mock_cache.return_value.delete.call_count == 1 - mock_lb.assert_called_once() - - session = Mock() - credential_record = SimpleNamespace(id="cred-1") - lb_config = SimpleNamespace(id="lb-1") - provider_model_record = SimpleNamespace(id="model-1", credential_id="cred-1", updated_at=None) - session.execute.side_effect = [ - _exec_result(scalar_one_or_none=credential_record), - _exec_result(scalars_all=[lb_config]), - _exec_result(scalar=2), - ] - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_custom_model_record", return_value=provider_model_record): - with patch("core.entities.provider_configuration.ProviderCredentialsCache") as mock_cache: - configuration.delete_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") - assert provider_model_record.credential_id is None - assert mock_cache.return_value.delete.call_count == 2 - assert mock_cache.call_args_list[0].kwargs["cache_type"] == ProviderCredentialsCacheType.LOAD_BALANCING_MODEL - assert mock_cache.call_args_list[1].kwargs["cache_type"] == ProviderCredentialsCacheType.MODEL - - session = Mock() - mismatched_credential_record = SimpleNamespace( - id="cred-2", - model_name="stored-model", - model_type=ModelType.TEXT_EMBEDDING, - ) - provider_model_record = SimpleNamespace(id="model-2", credential_id="cred-2", updated_at=None) - session.execute.side_effect = [ - _exec_result(scalar_one_or_none=None), - _exec_result(scalar_one_or_none=mismatched_credential_record), - _exec_result(scalars_all=[]), - _exec_result(scalar=1), - ] - with _patched_session(session): - with patch.object( - ProviderConfiguration, - "_get_custom_model_record", - return_value=provider_model_record, - ) as mock_get_model: - configuration.delete_custom_model_credential(ModelType.LLM, "request-model", "cred-2") - mock_get_model.assert_called_once_with(ModelType.TEXT_EMBEDDING, "stored-model", session=session) - session.delete.assert_any_call(mismatched_credential_record) - - -def test_add_model_credential_to_model_and_switch_custom_model_credential() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = None - with _patched_session(session): - with pytest.raises(ValueError, match="Credential record not found"): - configuration.add_model_credential_to_model(ModelType.LLM, "gpt-4o", "cred-1") - - session = Mock() - credential_record = SimpleNamespace(id="cred-1") - session.execute.return_value.scalar_one_or_none.return_value = credential_record - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_custom_model_record", return_value=None): - configuration.add_model_credential_to_model(ModelType.LLM, "gpt-4o", "cred-1") - session.add.assert_called_once() - session.commit.assert_called_once() - - session = Mock() - credential_record = SimpleNamespace(id="cred-1") - provider_model_record = SimpleNamespace(id="model-1", credential_id="cred-1", updated_at=None) - session.execute.return_value.scalar_one_or_none.return_value = credential_record - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_custom_model_record", return_value=provider_model_record): - with pytest.raises(ValueError, match="Can't add same credential"): - configuration.add_model_credential_to_model(ModelType.LLM, "gpt-4o", "cred-1") - - session = Mock() - credential_record = SimpleNamespace(id="cred-2") - provider_model_record = SimpleNamespace(id="model-1", credential_id="cred-1", updated_at=None) - session.execute.return_value.scalar_one_or_none.return_value = credential_record - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_custom_model_record", return_value=provider_model_record): - with patch("core.entities.provider_configuration.ProviderCredentialsCache") as mock_cache: - configuration.add_model_credential_to_model(ModelType.LLM, "gpt-4o", "cred-2") - assert provider_model_record.credential_id == "cred-2" - mock_cache.return_value.delete.assert_called_once() - - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = None - with _patched_session(session): - with pytest.raises(ValueError, match="Credential record not found"): - configuration.switch_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") - - session = Mock() - credential_record = SimpleNamespace(id="cred-1") - session.execute.return_value.scalar_one_or_none.return_value = credential_record - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_custom_model_record", return_value=None): - with pytest.raises(ValueError, match="custom model record not found"): - configuration.switch_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") - - session = Mock() - credential_record = SimpleNamespace(id="cred-1") - provider_model_record = SimpleNamespace(id="model-1", credential_id=None, updated_at=None) - session.execute.return_value.scalar_one_or_none.return_value = credential_record - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_custom_model_record", return_value=provider_model_record): - with patch("core.entities.provider_configuration.ProviderCredentialsCache") as mock_cache: - configuration.switch_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") - assert provider_model_record.credential_id == "cred-1" - mock_cache.return_value.delete.assert_called_once() - - -def test_delete_custom_model_and_model_setting_methods() -> None: - configuration = _build_provider_configuration() - session = Mock() - provider_model_record = SimpleNamespace(id="model-1") - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_custom_model_record", return_value=provider_model_record): - with patch("core.entities.provider_configuration.ProviderCredentialsCache") as mock_cache: - configuration.delete_custom_model(ModelType.LLM, "gpt-4o") - session.delete.assert_called_once_with(provider_model_record) - session.commit.assert_called_once() - mock_cache.return_value.delete.assert_called_once() - - session = Mock() - existing = SimpleNamespace(enabled=False, updated_at=None) - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_model_setting", return_value=existing): - assert configuration.enable_model(ModelType.LLM, "gpt-4o") is existing - assert existing.enabled is True - - session = Mock() - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_model_setting", return_value=None): - created = configuration.enable_model(ModelType.LLM, "gpt-4o") - assert created.enabled is True - - session = Mock() - existing = SimpleNamespace(enabled=True, load_balancing_enabled=True, updated_at=None) - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_model_setting", return_value=existing): - assert configuration.disable_model(ModelType.LLM, "gpt-4o") is existing - assert existing.enabled is False - - session = Mock() - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_model_setting", return_value=None): - created = configuration.disable_model(ModelType.LLM, "gpt-4o") - assert created.enabled is False - - session = Mock() - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_model_setting", return_value=existing): - result = configuration.get_provider_model_setting(ModelType.LLM, "gpt-4o") - assert result is existing - - -def test_model_load_balancing_enable_disable_and_switch_preferred_provider_type_without_session() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.execute.return_value.scalar.return_value = 1 - with _patched_session(session): - with pytest.raises(ValueError, match="must be more than 1"): - configuration.enable_model_load_balancing(ModelType.LLM, "gpt-4o") - - session = Mock() - session.execute.return_value.scalar.return_value = 2 - existing = SimpleNamespace(load_balancing_enabled=False, updated_at=None) - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_model_setting", return_value=existing): - result = configuration.enable_model_load_balancing(ModelType.LLM, "gpt-4o") - assert result is existing - assert existing.load_balancing_enabled is True - - session = Mock() - session.execute.return_value.scalar.return_value = 2 - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_model_setting", return_value=None): - created = configuration.enable_model_load_balancing(ModelType.LLM, "gpt-4o") - assert created.load_balancing_enabled is True - - session = Mock() - existing = SimpleNamespace(load_balancing_enabled=True, updated_at=None) - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_model_setting", return_value=existing): - result = configuration.disable_model_load_balancing(ModelType.LLM, "gpt-4o") - assert result is existing - assert existing.load_balancing_enabled is False - - session = Mock() - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_model_setting", return_value=None): - created = configuration.disable_model_load_balancing(ModelType.LLM, "gpt-4o") - assert created.load_balancing_enabled is False - - configuration.preferred_provider_type = ProviderType.SYSTEM - switch_session = Mock() - with _patched_session(switch_session): - switch_session.execute.return_value.scalars.return_value.first.return_value = None - with patch.object(ProviderConfiguration, "_invalidate_provider_configuration_cache") as mock_invalidate: - changed = configuration.switch_preferred_provider_type(ProviderType.CUSTOM) - assert changed is True - assert any( - call.args and call.args[0].__class__.__name__ == "TenantPreferredModelProvider" - for call in switch_session.add.call_args_list - ) - switch_session.commit.assert_called() - mock_invalidate.assert_called_once_with(preferred_model_providers=True) - - -def test_system_and_custom_provider_model_helpers_cover_remaining_skip_paths() -> None: - configuration = _build_provider_configuration() - provider_schema = ProviderEntity( - provider="openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.CUSTOMIZABLE_MODEL], - models=[_build_ai_model("llm-model")], - ) - configuration.system_configuration.quota_configurations = [ - QuotaConfiguration( - quota_type=ProviderQuotaType.FREE, - quota_unit=QuotaUnit.TOKENS, - quota_limit=100, - quota_used=0, - is_valid=True, - restrict_models=[ - RestrictModel(model="target", base_model_name="base", model_type=ModelType.LLM), - ], - ), - QuotaConfiguration( - quota_type=ProviderQuotaType.TRIAL, - quota_unit=QuotaUnit.TOKENS, - quota_limit=100, - quota_used=0, - is_valid=True, - restrict_models=[ - RestrictModel(model="target", base_model_name="base", model_type=ModelType.LLM), - RestrictModel(model="error-model", base_model_name="base", model_type=ModelType.LLM), - RestrictModel(model="none-model", base_model_name="base", model_type=ModelType.LLM), - RestrictModel( - model="embed-model", - base_model_name="base", - model_type=ModelType.TEXT_EMBEDDING, - ), - ], - ), - ] - configuration.system_configuration.current_quota_type = ProviderQuotaType.TRIAL - - def _system_schema(*, model_type: ModelType, model: str, credentials: dict | None): - if model == "error-model": - raise RuntimeError("boom") - if model == "none-model": - return None - if model == "embed-model": - return _build_ai_model("embed-model", model_type=ModelType.TEXT_EMBEDDING) - return _build_ai_model("target") - - configuration._original_provider_configurate_methods = (ConfigurateMethod.CUSTOMIZABLE_MODEL,) - with patch.object(ProviderConfiguration, "get_model_schema", side_effect=_system_schema): - system_models = configuration._get_system_provider_models( - model_types=[ModelType.LLM], - provider_schema=provider_schema, - model_setting_map={ - ModelType.LLM: { - "target": ModelSettings( - model="target", - model_type=ModelType.LLM, - enabled=False, - load_balancing_configs=[], - ) - } - }, - ) - assert any(model.model == "target" and model.status == ModelStatus.DISABLED for model in system_models) - - configuration.using_provider_type = ProviderType.CUSTOM - configuration.custom_configuration.provider = CustomProviderConfiguration(credentials={"api_key": "provider-key"}) - configuration.custom_configuration.models = [ - CustomModelConfiguration( - model="skip-model-type", - model_type=ModelType.TEXT_EMBEDDING, - credentials={"k": "v"}, - ), - CustomModelConfiguration( - model="skip-unadded", - model_type=ModelType.LLM, - credentials={"k": "v"}, - unadded_to_model_list=True, - ), - CustomModelConfiguration( - model="skip-filter", - model_type=ModelType.LLM, - credentials={"k": "v"}, - ), - CustomModelConfiguration( - model="error-custom", - model_type=ModelType.LLM, - credentials={"k": "v"}, - ), - CustomModelConfiguration( - model="none-custom", - model_type=ModelType.LLM, - credentials={"k": "v"}, - ), - CustomModelConfiguration( - model="disabled-custom", - model_type=ModelType.LLM, - credentials={"k": "v"}, - ), - ] - - provider_schema = ProviderEntity( - provider="openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - models=[_build_ai_model("base-disabled")], - ) - model_setting_map = { - ModelType.LLM: { - "base-disabled": ModelSettings( - model="base-disabled", - model_type=ModelType.LLM, - enabled=False, - load_balancing_enabled=True, - load_balancing_configs=[ModelLoadBalancingConfiguration(id="lb-1", name="lb", credentials={})], - ), - "disabled-custom": ModelSettings( - model="disabled-custom", - model_type=ModelType.LLM, - enabled=False, - load_balancing_enabled=False, - load_balancing_configs=[], - ), - } - } - - def _custom_schema(*, model_type: ModelType, model: str, credentials: dict | None): - if model == "error-custom": - raise RuntimeError("boom") - if model == "none-custom": - return None - return _build_ai_model(model) - - with patch.object(ProviderConfiguration, "get_model_schema", side_effect=_custom_schema): - custom_models = configuration._get_custom_provider_models( - model_types=[ModelType.LLM], - provider_schema=provider_schema, - model_setting_map=model_setting_map, - model="disabled-custom", - ) - assert any(model.model == "base-disabled" and model.status == ModelStatus.DISABLED for model in custom_models) - assert any(model.model == "disabled-custom" and model.status == ModelStatus.DISABLED for model in custom_models) - - -def test_get_current_credentials_skips_non_current_quota_restrictions() -> None: - configuration = _build_provider_configuration() - configuration.system_configuration.current_quota_type = ProviderQuotaType.TRIAL - configuration.system_configuration.quota_configurations = [ - QuotaConfiguration( - quota_type=ProviderQuotaType.FREE, - quota_unit=QuotaUnit.TOKENS, - quota_limit=100, - quota_used=0, - is_valid=True, - restrict_models=[ - RestrictModel(model="gpt-4o", base_model_name="free-base", model_type=ModelType.LLM), - ], - ), - QuotaConfiguration( - quota_type=ProviderQuotaType.TRIAL, - quota_unit=QuotaUnit.TOKENS, - quota_limit=100, - quota_used=0, - is_valid=True, - restrict_models=[ - RestrictModel(model="gpt-4o", base_model_name="trial-base", model_type=ModelType.LLM), - ], - ), - ] - - credentials = configuration.get_current_credentials(ModelType.LLM, "gpt-4o") - assert credentials["base_model_name"] == "trial-base" - - -def test_get_system_configuration_status_covers_disabled_and_quota_exceeded() -> None: - configuration = _build_provider_configuration() - configuration.system_configuration.enabled = False - assert configuration.get_system_configuration_status() == SystemConfigurationStatus.UNSUPPORTED - - configuration.system_configuration.enabled = True - configuration.system_configuration.quota_configurations = [ - QuotaConfiguration( - quota_type=ProviderQuotaType.TRIAL, - quota_unit=QuotaUnit.TOKENS, - quota_limit=100, - quota_used=100, - is_valid=False, - restrict_models=[], - ) - ] - configuration.system_configuration.current_quota_type = ProviderQuotaType.TRIAL - assert configuration.get_system_configuration_status() == SystemConfigurationStatus.QUOTA_EXCEEDED - - -def test_get_specific_provider_credential_decrypts_and_obfuscates_credentials() -> None: - configuration = _build_provider_configuration() - configuration.provider.provider_credential_schema = _build_secret_provider_schema() - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = SimpleNamespace( - encrypted_config='{"openai_api_key":"enc-secret","region":"us"}' - ) - provider_record = SimpleNamespace(provider_name="aliased-openai") - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_record", return_value=provider_record): - with patch("core.entities.provider_configuration.encrypter.decrypt_token", return_value="raw-secret"): - with patch.object( - ProviderConfiguration, - "obfuscated_credentials", - side_effect=lambda credentials, credential_form_schemas: credentials, - ): - credentials = configuration._get_specific_provider_credential("cred-1") - - assert credentials == {"openai_api_key": "raw-secret", "region": "us"} - - -def test_get_specific_provider_credential_logs_when_decrypt_fails(caplog: pytest.LogCaptureFixture) -> None: - configuration = _build_provider_configuration() - configuration.provider.provider_credential_schema = _build_secret_provider_schema() - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = SimpleNamespace( - encrypted_config='{"openai_api_key":"enc-secret"}' - ) - caplog.set_level(logging.ERROR, logger="core.entities.provider_configuration") - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_record", return_value=None): - with patch( - "core.entities.provider_configuration.encrypter.decrypt_token", - side_effect=RuntimeError("boom"), - ): - with patch.object( - ProviderConfiguration, - "obfuscated_credentials", - side_effect=lambda credentials, credential_form_schemas: credentials, - ): - credentials = configuration._get_specific_provider_credential("cred-1") - - assert credentials == {"openai_api_key": "enc-secret"} - assert caplog.messages.count("Failed to decrypt credential secret variable openai_api_key") == 1 - - -def test_validate_provider_credentials_uses_empty_original_when_record_missing() -> None: - configuration = _build_provider_configuration() - configuration.provider.provider_credential_schema = _build_secret_provider_schema() - mock_session = Mock() - mock_session.execute.return_value.scalar_one_or_none.return_value = None - mock_factory = Mock() - mock_factory.provider_credentials_validate.return_value = {"openai_api_key": "raw"} - - with _patched_session(mock_session): - with patch( - "core.entities.provider_configuration.create_plugin_model_assembly", - return_value=SimpleNamespace(model_runtime=Mock(), model_provider_factory=mock_factory), - ): - with patch("core.entities.provider_configuration.encrypter.encrypt_token", return_value="enc-new"): - validated = configuration.validate_provider_credentials( - credentials={"openai_api_key": HIDDEN_VALUE}, - credential_id="cred-1", - ) - - assert validated == {"openai_api_key": "enc-new"} - - -def test_create_provider_credential_rolls_back_on_error() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.add.side_effect = RuntimeError("boom") - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_generate_provider_credential_name", return_value="API KEY 9"): - with patch.object(ProviderConfiguration, "validate_provider_credentials", return_value={"api_key": "enc"}): - with patch.object(ProviderConfiguration, "_get_provider_record", return_value=None): - with pytest.raises(RuntimeError, match="boom"): - configuration.create_provider_credential({"api_key": "raw"}, None) - - session.rollback.assert_called_once() - - -def test_update_provider_credential_raises_on_duplicate_name() -> None: - configuration = _build_provider_configuration() - session = Mock() - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_check_provider_credential_name_exists", return_value=True): - with pytest.raises(ValueError, match="already exists"): - configuration.update_provider_credential({"api_key": "raw"}, "cred-1", "Main") - - -def test_update_provider_credential_rolls_back_on_error() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = SimpleNamespace( - id="cred-1", - encrypted_config="{}", - credential_name="Main", - updated_at=None, - ) - session.commit.side_effect = RuntimeError("boom") - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_check_provider_credential_name_exists", return_value=False): - with patch.object(ProviderConfiguration, "validate_provider_credentials", return_value={"api_key": "enc"}): - with patch.object(ProviderConfiguration, "_get_provider_record", return_value=None): - with pytest.raises(RuntimeError, match="boom"): - configuration.update_provider_credential({"api_key": "raw"}, "cred-1", "Main") - - session.rollback.assert_called_once() - - -def test_delete_provider_credential_rolls_back_on_error() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.delete.side_effect = RuntimeError("boom") - session.execute.side_effect = [ - _exec_result(scalar_one_or_none=SimpleNamespace(id="cred-1")), - _exec_result(scalars_all=[]), - _exec_result(scalar=2), - ] - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_record", return_value=None): - with pytest.raises(RuntimeError, match="boom"): - configuration.delete_provider_credential("cred-1") - - session.rollback.assert_called_once() - - -def test_switch_active_provider_credential_rolls_back_on_error() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = SimpleNamespace(id="cred-1") - session.commit.side_effect = RuntimeError("boom") - provider_record = SimpleNamespace(id="provider-1", credential_id=None, updated_at=None) - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_provider_record", return_value=provider_record): - with pytest.raises(RuntimeError, match="boom"): - configuration.switch_active_provider_credential("cred-1") - - session.rollback.assert_called_once() - - -def test_get_specific_custom_model_credential_logs_when_decrypt_fails(caplog: pytest.LogCaptureFixture) -> None: - configuration = _build_provider_configuration() - configuration.provider.model_credential_schema = _build_secret_model_schema() - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = SimpleNamespace( - id="cred-1", - credential_name="Main", - encrypted_config='{"openai_api_key":"enc-secret"}', - ) - caplog.set_level(logging.ERROR, logger="core.entities.provider_configuration") - - with _patched_session(session): - with patch("core.entities.provider_configuration.encrypter.decrypt_token", side_effect=RuntimeError("boom")): - with patch.object( - ProviderConfiguration, - "obfuscated_credentials", - side_effect=lambda credentials, credential_form_schemas: credentials, - ): - result = configuration._get_specific_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") - - assert result["credentials"] == {"openai_api_key": "enc-secret"} - assert caplog.messages.count("Failed to decrypt model credential secret variable openai_api_key") == 1 - - -def test_validate_custom_model_credentials_handles_invalid_original_json() -> None: - configuration = _build_provider_configuration() - configuration.provider.model_credential_schema = _build_secret_model_schema() - mock_session = Mock() - mock_session.execute.return_value.scalar_one_or_none.return_value = SimpleNamespace( - encrypted_config="{invalid-json" - ) - mock_factory = Mock() - mock_factory.model_credentials_validate.return_value = {"openai_api_key": "raw"} - - with _patched_session(mock_session): - with patch( - "core.entities.provider_configuration.create_plugin_model_assembly", - return_value=SimpleNamespace(model_runtime=Mock(), model_provider_factory=mock_factory), - ): - with patch("core.entities.provider_configuration.encrypter.encrypt_token", return_value="enc-new"): - validated = configuration.validate_custom_model_credentials( - model_type=ModelType.LLM, - model="gpt-4o", - credentials={"openai_api_key": HIDDEN_VALUE}, - credential_id="cred-1", - ) - - assert validated == {"openai_api_key": "enc-new"} - - -def test_create_custom_model_credential_raises_on_duplicate_name() -> None: - configuration = _build_provider_configuration() - session = Mock() - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_check_custom_model_credential_name_exists", return_value=True): - with pytest.raises(ValueError, match="already exists"): - configuration.create_custom_model_credential(ModelType.LLM, "gpt-4o", {"k": "v"}, "Main") - - -def test_create_custom_model_credential_rolls_back_on_error() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.add.side_effect = RuntimeError("boom") - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_generate_custom_model_credential_name", return_value="API KEY 4"): - with patch.object( - ProviderConfiguration, - "validate_custom_model_credentials", - return_value={"openai_api_key": "enc"}, - ): - with patch.object(ProviderConfiguration, "_get_custom_model_record", return_value=None): - with pytest.raises(RuntimeError, match="boom"): - configuration.create_custom_model_credential(ModelType.LLM, "gpt-4o", {"k": "v"}, None) - - session.rollback.assert_called_once() - - -def test_update_custom_model_credential_raises_on_duplicate_name() -> None: - configuration = _build_provider_configuration() - session = Mock() - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_check_custom_model_credential_name_exists", return_value=True): - with pytest.raises(ValueError, match="already exists"): - configuration.update_custom_model_credential( - model_type=ModelType.LLM, - model="gpt-4o", - credentials={"k": "v"}, - credential_name="Main", - credential_id="cred-1", - ) - - -def test_update_custom_model_credential_raises_when_record_not_found() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = None - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_check_custom_model_credential_name_exists", return_value=False): - with patch.object(ProviderConfiguration, "validate_custom_model_credentials", return_value={"k": "v"}): - with patch.object(ProviderConfiguration, "_get_custom_model_record", return_value=None): - with pytest.raises(ValueError, match="Credential record not found"): - configuration.update_custom_model_credential( - model_type=ModelType.LLM, - model="gpt-4o", - credentials={"k": "v"}, - credential_name="Main", - credential_id="cred-1", - ) - - -def test_update_custom_model_credential_rolls_back_on_error() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = SimpleNamespace( - id="cred-1", - encrypted_config="{}", - credential_name="Main", - updated_at=None, - ) - session.commit.side_effect = RuntimeError("boom") - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_check_custom_model_credential_name_exists", return_value=False): - with patch.object(ProviderConfiguration, "validate_custom_model_credentials", return_value={"k": "v"}): - with patch.object(ProviderConfiguration, "_get_custom_model_record", return_value=None): - with pytest.raises(RuntimeError, match="boom"): - configuration.update_custom_model_credential( - model_type=ModelType.LLM, - model="gpt-4o", - credentials={"k": "v"}, - credential_name="Main", - credential_id="cred-1", - ) - - session.rollback.assert_called_once() - - -def test_delete_custom_model_credential_raises_when_record_not_found() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.execute.return_value.scalar_one_or_none.return_value = None - - with _patched_session(session): - with pytest.raises(ValueError, match="Credential record not found"): - configuration.delete_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") - - -def test_delete_custom_model_credential_removes_custom_model_record_when_last_credential() -> None: - configuration = _build_provider_configuration() - session = Mock() - credential_record = SimpleNamespace(id="cred-1") - provider_model_record = SimpleNamespace(id="model-1", credential_id="cred-1", updated_at=None) - session.execute.side_effect = [ - _exec_result(scalar_one_or_none=credential_record), - _exec_result(scalars_all=[]), - _exec_result(scalar=1), - ] - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_custom_model_record", return_value=provider_model_record): - with patch("core.entities.provider_configuration.ProviderCredentialsCache") as mock_cache: - configuration.delete_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") - - assert any(call.args and call.args[0] is provider_model_record for call in session.delete.call_args_list) - mock_cache.assert_called_once_with( - tenant_id="tenant-1", - identity_id="model-1", - cache_type=ProviderCredentialsCacheType.MODEL, - ) - mock_cache.return_value.delete.assert_called_once() - - -def test_delete_custom_model_credential_rolls_back_on_error() -> None: - configuration = _build_provider_configuration() - session = Mock() - session.delete.side_effect = RuntimeError("boom") - session.execute.side_effect = [ - _exec_result(scalar_one_or_none=SimpleNamespace(id="cred-1")), - _exec_result(scalars_all=[]), - _exec_result(scalar=2), - ] - - with _patched_session(session): - with patch.object(ProviderConfiguration, "_get_custom_model_record", return_value=None): - with pytest.raises(RuntimeError, match="boom"): - configuration.delete_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") - - session.rollback.assert_called_once() - - def test_get_custom_provider_models_skips_schema_models_with_mismatched_type() -> None: configuration = _build_provider_configuration() provider_schema = ProviderEntity(