From 24cd7bbc622a9bb4ae8384e51ab6722ca3e92210 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 13 Oct 2025 11:29:37 +0900 Subject: [PATCH] fix RetrievalMethod StrEnum (#26768) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: crazywoola <100913391+crazywoola@users.noreply.github.com> --- .../entities/datasource_entities.py | 4 ++-- .../entities/provider_entities.py | 4 ++-- api/core/rag/datasource/retrieval_service.py | 6 +++--- api/core/rag/entities/event.py | 4 ++-- .../index_processor/index_processor_base.py | 3 ++- .../processor/paragraph_index_processor.py | 3 ++- .../processor/parent_child_index_processor.py | 3 ++- .../processor/qa_index_processor.py | 3 ++- api/core/rag/retrieval/dataset_retrieval.py | 4 ++-- api/core/rag/retrieval/retrieval_methods.py | 4 ++-- .../dataset_multi_retriever_tool.py | 2 +- .../dataset_retriever_tool.py | 2 +- api/core/workflow/enums.py | 4 ++-- .../graph_engine/layers/execution_limits.py | 4 ++-- .../nodes/knowledge_index/entities.py | 3 ++- api/services/dataset_service.py | 6 +++--- .../knowledge_entities/knowledge_entities.py | 4 +++- .../rag_pipeline_entities.py | 4 +++- .../entities/model_provider_entities.py | 4 ++-- api/services/hit_testing_service.py | 2 +- .../rag_pipeline_transform_service.py | 3 ++- .../rag/extractor/firecrawl/test_firecrawl.py | 4 +++- .../rag/extractor/test_notion_extractor.py | 6 ++++-- .../unit_tests/core/test_model_manager.py | 3 ++- .../unit_tests/core/test_provider_manager.py | 19 +++++++++++++------ 25 files changed, 65 insertions(+), 43 deletions(-) diff --git a/api/core/datasource/entities/datasource_entities.py b/api/core/datasource/entities/datasource_entities.py index 7f503b963f..260dcf04f5 100644 --- a/api/core/datasource/entities/datasource_entities.py +++ b/api/core/datasource/entities/datasource_entities.py @@ -1,5 +1,5 @@ import enum -from enum import Enum +from enum import StrEnum from typing import Any from pydantic import BaseModel, Field, ValidationInfo, field_validator @@ -218,7 +218,7 @@ class DatasourceLabel(BaseModel): icon: str = Field(..., description="The icon of the tool") -class DatasourceInvokeFrom(Enum): +class DatasourceInvokeFrom(StrEnum): """ Enum class for datasource invoke """ diff --git a/api/core/model_runtime/entities/provider_entities.py b/api/core/model_runtime/entities/provider_entities.py index 831fb9d4db..0508116962 100644 --- a/api/core/model_runtime/entities/provider_entities.py +++ b/api/core/model_runtime/entities/provider_entities.py @@ -1,5 +1,5 @@ from collections.abc import Sequence -from enum import Enum, StrEnum, auto +from enum import StrEnum, auto from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator @@ -7,7 +7,7 @@ from core.model_runtime.entities.common_entities import I18nObject from core.model_runtime.entities.model_entities import AIModelEntity, ModelType -class ConfigurateMethod(Enum): +class ConfigurateMethod(StrEnum): """ Enum class for configurate method of provider model. """ diff --git a/api/core/rag/datasource/retrieval_service.py b/api/core/rag/datasource/retrieval_service.py index 6e9e2b4527..2290de19bc 100644 --- a/api/core/rag/datasource/retrieval_service.py +++ b/api/core/rag/datasource/retrieval_service.py @@ -34,7 +34,7 @@ class RetrievalService: @classmethod def retrieve( cls, - retrieval_method: str, + retrieval_method: RetrievalMethod, dataset_id: str, query: str, top_k: int, @@ -56,7 +56,7 @@ class RetrievalService: # Optimize multithreading with thread pools with ThreadPoolExecutor(max_workers=dify_config.RETRIEVAL_SERVICE_EXECUTORS) as executor: # type: ignore futures = [] - if retrieval_method == "keyword_search": + if retrieval_method == RetrievalMethod.KEYWORD_SEARCH: futures.append( executor.submit( cls.keyword_search, @@ -220,7 +220,7 @@ class RetrievalService: score_threshold: float | None, reranking_model: dict | None, all_documents: list, - retrieval_method: str, + retrieval_method: RetrievalMethod, exceptions: list, document_ids_filter: list[str] | None = None, ): diff --git a/api/core/rag/entities/event.py b/api/core/rag/entities/event.py index a61b17ddb8..2d8d4060dd 100644 --- a/api/core/rag/entities/event.py +++ b/api/core/rag/entities/event.py @@ -1,11 +1,11 @@ from collections.abc import Mapping -from enum import Enum +from enum import StrEnum from typing import Any from pydantic import BaseModel, Field -class DatasourceStreamEvent(Enum): +class DatasourceStreamEvent(StrEnum): """ Datasource Stream event """ diff --git a/api/core/rag/index_processor/index_processor_base.py b/api/core/rag/index_processor/index_processor_base.py index 05cffb5a55..d4eff53204 100644 --- a/api/core/rag/index_processor/index_processor_base.py +++ b/api/core/rag/index_processor/index_processor_base.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Optional from configs import dify_config from core.rag.extractor.entity.extract_setting import ExtractSetting from core.rag.models.document import Document +from core.rag.retrieval.retrieval_methods import RetrievalMethod from core.rag.splitter.fixed_text_splitter import ( EnhanceRecursiveCharacterTextSplitter, FixedRecursiveCharacterTextSplitter, @@ -49,7 +50,7 @@ class BaseIndexProcessor(ABC): @abstractmethod def retrieve( self, - retrieval_method: str, + retrieval_method: RetrievalMethod, query: str, dataset: Dataset, top_k: int, diff --git a/api/core/rag/index_processor/processor/paragraph_index_processor.py b/api/core/rag/index_processor/processor/paragraph_index_processor.py index 4fcffbcc77..5e5fea7ea9 100644 --- a/api/core/rag/index_processor/processor/paragraph_index_processor.py +++ b/api/core/rag/index_processor/processor/paragraph_index_processor.py @@ -14,6 +14,7 @@ from core.rag.extractor.extract_processor import ExtractProcessor from core.rag.index_processor.constant.index_type import IndexType from core.rag.index_processor.index_processor_base import BaseIndexProcessor from core.rag.models.document import Document +from core.rag.retrieval.retrieval_methods import RetrievalMethod from core.tools.utils.text_processing_utils import remove_leading_symbols from libs import helper from models.dataset import Dataset, DatasetProcessRule @@ -106,7 +107,7 @@ class ParagraphIndexProcessor(BaseIndexProcessor): def retrieve( self, - retrieval_method: str, + retrieval_method: RetrievalMethod, query: str, dataset: Dataset, top_k: int, diff --git a/api/core/rag/index_processor/processor/parent_child_index_processor.py b/api/core/rag/index_processor/processor/parent_child_index_processor.py index 7bdde286f5..4fa78e2f95 100644 --- a/api/core/rag/index_processor/processor/parent_child_index_processor.py +++ b/api/core/rag/index_processor/processor/parent_child_index_processor.py @@ -16,6 +16,7 @@ from core.rag.extractor.extract_processor import ExtractProcessor from core.rag.index_processor.constant.index_type import IndexType from core.rag.index_processor.index_processor_base import BaseIndexProcessor from core.rag.models.document import ChildDocument, Document, ParentChildStructureChunk +from core.rag.retrieval.retrieval_methods import RetrievalMethod from extensions.ext_database import db from libs import helper from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentSegment @@ -161,7 +162,7 @@ class ParentChildIndexProcessor(BaseIndexProcessor): def retrieve( self, - retrieval_method: str, + retrieval_method: RetrievalMethod, query: str, dataset: Dataset, top_k: int, diff --git a/api/core/rag/index_processor/processor/qa_index_processor.py b/api/core/rag/index_processor/processor/qa_index_processor.py index 9c8f70dba8..3e3deb0180 100644 --- a/api/core/rag/index_processor/processor/qa_index_processor.py +++ b/api/core/rag/index_processor/processor/qa_index_processor.py @@ -21,6 +21,7 @@ from core.rag.extractor.extract_processor import ExtractProcessor from core.rag.index_processor.constant.index_type import IndexType from core.rag.index_processor.index_processor_base import BaseIndexProcessor from core.rag.models.document import Document, QAStructureChunk +from core.rag.retrieval.retrieval_methods import RetrievalMethod from core.tools.utils.text_processing_utils import remove_leading_symbols from libs import helper from models.dataset import Dataset @@ -141,7 +142,7 @@ class QAIndexProcessor(BaseIndexProcessor): def retrieve( self, - retrieval_method: str, + retrieval_method: RetrievalMethod, query: str, dataset: Dataset, top_k: int, diff --git a/api/core/rag/retrieval/dataset_retrieval.py b/api/core/rag/retrieval/dataset_retrieval.py index 0a702d2902..99bbe615fb 100644 --- a/api/core/rag/retrieval/dataset_retrieval.py +++ b/api/core/rag/retrieval/dataset_retrieval.py @@ -364,7 +364,7 @@ class DatasetRetrieval: top_k = retrieval_model_config["top_k"] # get retrieval method if dataset.indexing_technique == "economy": - retrieval_method = "keyword_search" + retrieval_method = RetrievalMethod.KEYWORD_SEARCH else: retrieval_method = retrieval_model_config["search_method"] # get reranking model @@ -623,7 +623,7 @@ class DatasetRetrieval: if dataset.indexing_technique == "economy": # use keyword table query documents = RetrievalService.retrieve( - retrieval_method="keyword_search", + retrieval_method=RetrievalMethod.KEYWORD_SEARCH, dataset_id=dataset.id, query=query, top_k=top_k, diff --git a/api/core/rag/retrieval/retrieval_methods.py b/api/core/rag/retrieval/retrieval_methods.py index 5f0f2a9d33..c77a026351 100644 --- a/api/core/rag/retrieval/retrieval_methods.py +++ b/api/core/rag/retrieval/retrieval_methods.py @@ -1,7 +1,7 @@ -from enum import Enum +from enum import StrEnum -class RetrievalMethod(Enum): +class RetrievalMethod(StrEnum): SEMANTIC_SEARCH = "semantic_search" FULL_TEXT_SEARCH = "full_text_search" HYBRID_SEARCH = "hybrid_search" diff --git a/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py b/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py index b5bc4d3c00..20e10be075 100644 --- a/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py +++ b/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py @@ -172,7 +172,7 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool): if dataset.indexing_technique == "economy": # use keyword table query documents = RetrievalService.retrieve( - retrieval_method="keyword_search", + retrieval_method=RetrievalMethod.KEYWORD_SEARCH, dataset_id=dataset.id, query=query, top_k=retrieval_model.get("top_k") or 4, diff --git a/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py b/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py index 1eae582f67..915a22dd0f 100644 --- a/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py +++ b/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py @@ -130,7 +130,7 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool): if dataset.indexing_technique == "economy": # use keyword table query documents = RetrievalService.retrieve( - retrieval_method="keyword_search", + retrieval_method=RetrievalMethod.KEYWORD_SEARCH, dataset_id=dataset.id, query=query, top_k=self.top_k, diff --git a/api/core/workflow/enums.py b/api/core/workflow/enums.py index 00a125660a..eb88bb67ee 100644 --- a/api/core/workflow/enums.py +++ b/api/core/workflow/enums.py @@ -1,7 +1,7 @@ -from enum import Enum, StrEnum +from enum import StrEnum -class NodeState(Enum): +class NodeState(StrEnum): """State of a node or edge during workflow execution.""" UNKNOWN = "unknown" diff --git a/api/core/workflow/graph_engine/layers/execution_limits.py b/api/core/workflow/graph_engine/layers/execution_limits.py index e39af89837..a2d36d142d 100644 --- a/api/core/workflow/graph_engine/layers/execution_limits.py +++ b/api/core/workflow/graph_engine/layers/execution_limits.py @@ -10,7 +10,7 @@ When limits are exceeded, the layer automatically aborts execution. import logging import time -from enum import Enum +from enum import StrEnum from typing import final from typing_extensions import override @@ -24,7 +24,7 @@ from core.workflow.graph_events import ( from core.workflow.graph_events.node import NodeRunFailedEvent, NodeRunSucceededEvent -class LimitType(Enum): +class LimitType(StrEnum): """Types of execution limits that can be exceeded.""" STEP_LIMIT = "step_limit" diff --git a/api/core/workflow/nodes/knowledge_index/entities.py b/api/core/workflow/nodes/knowledge_index/entities.py index c79373afd5..3daca90b9b 100644 --- a/api/core/workflow/nodes/knowledge_index/entities.py +++ b/api/core/workflow/nodes/knowledge_index/entities.py @@ -2,6 +2,7 @@ from typing import Literal, Union from pydantic import BaseModel +from core.rag.retrieval.retrieval_methods import RetrievalMethod from core.workflow.nodes.base import BaseNodeData @@ -63,7 +64,7 @@ class RetrievalSetting(BaseModel): Retrieval Setting. """ - search_method: Literal["semantic_search", "keyword_search", "full_text_search", "hybrid_search"] + search_method: RetrievalMethod top_k: int score_threshold: float | None = 0.5 score_threshold_enabled: bool = False diff --git a/api/services/dataset_service.py b/api/services/dataset_service.py index 87861ada87..53216e4fdd 100644 --- a/api/services/dataset_service.py +++ b/api/services/dataset_service.py @@ -1470,7 +1470,7 @@ class DocumentService: dataset.collection_binding_id = dataset_collection_binding.id if not dataset.retrieval_model: default_retrieval_model = { - "search_method": RetrievalMethod.SEMANTIC_SEARCH.value, + "search_method": RetrievalMethod.SEMANTIC_SEARCH, "reranking_enable": False, "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""}, "top_k": 4, @@ -1752,7 +1752,7 @@ class DocumentService: # dataset.collection_binding_id = dataset_collection_binding.id # if not dataset.retrieval_model: # default_retrieval_model = { - # "search_method": RetrievalMethod.SEMANTIC_SEARCH.value, + # "search_method": RetrievalMethod.SEMANTIC_SEARCH, # "reranking_enable": False, # "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""}, # "top_k": 2, @@ -2205,7 +2205,7 @@ class DocumentService: retrieval_model = knowledge_config.retrieval_model else: retrieval_model = RetrievalModel( - search_method=RetrievalMethod.SEMANTIC_SEARCH.value, + search_method=RetrievalMethod.SEMANTIC_SEARCH, reranking_enable=False, reranking_model=RerankingModel(reranking_provider_name="", reranking_model_name=""), top_k=4, diff --git a/api/services/entities/knowledge_entities/knowledge_entities.py b/api/services/entities/knowledge_entities/knowledge_entities.py index 33f65bde58..b9a210740d 100644 --- a/api/services/entities/knowledge_entities/knowledge_entities.py +++ b/api/services/entities/knowledge_entities/knowledge_entities.py @@ -3,6 +3,8 @@ from typing import Literal from pydantic import BaseModel +from core.rag.retrieval.retrieval_methods import RetrievalMethod + class ParentMode(StrEnum): FULL_DOC = "full-doc" @@ -95,7 +97,7 @@ class WeightModel(BaseModel): class RetrievalModel(BaseModel): - search_method: Literal["hybrid_search", "semantic_search", "full_text_search", "keyword_search"] + search_method: RetrievalMethod reranking_enable: bool reranking_model: RerankingModel | None = None reranking_mode: str | None = None diff --git a/api/services/entities/knowledge_entities/rag_pipeline_entities.py b/api/services/entities/knowledge_entities/rag_pipeline_entities.py index 860bfde401..a97ccab914 100644 --- a/api/services/entities/knowledge_entities/rag_pipeline_entities.py +++ b/api/services/entities/knowledge_entities/rag_pipeline_entities.py @@ -2,6 +2,8 @@ from typing import Literal from pydantic import BaseModel, field_validator +from core.rag.retrieval.retrieval_methods import RetrievalMethod + class IconInfo(BaseModel): icon: str @@ -83,7 +85,7 @@ class RetrievalSetting(BaseModel): Retrieval Setting. """ - search_method: Literal["semantic_search", "full_text_search", "keyword_search", "hybrid_search"] + search_method: RetrievalMethod top_k: int score_threshold: float | None = 0.5 score_threshold_enabled: bool = False diff --git a/api/services/entities/model_provider_entities.py b/api/services/entities/model_provider_entities.py index 0f5151919f..d07badefa7 100644 --- a/api/services/entities/model_provider_entities.py +++ b/api/services/entities/model_provider_entities.py @@ -1,5 +1,5 @@ from collections.abc import Sequence -from enum import Enum +from enum import StrEnum from pydantic import BaseModel, ConfigDict, model_validator @@ -27,7 +27,7 @@ from core.model_runtime.entities.provider_entities import ( from models.provider import ProviderType -class CustomConfigurationStatus(Enum): +class CustomConfigurationStatus(StrEnum): """ Enum class for custom configuration status. """ diff --git a/api/services/hit_testing_service.py b/api/services/hit_testing_service.py index aa29354a6e..c6ea35076e 100644 --- a/api/services/hit_testing_service.py +++ b/api/services/hit_testing_service.py @@ -63,7 +63,7 @@ class HitTestingService: if metadata_condition and not document_ids_filter: return cls.compact_retrieve_response(query, []) all_documents = RetrievalService.retrieve( - retrieval_method=retrieval_model.get("search_method", "semantic_search"), + retrieval_method=RetrievalMethod(retrieval_model.get("search_method", RetrievalMethod.SEMANTIC_SEARCH)), dataset_id=dataset.id, query=query, top_k=retrieval_model.get("top_k", 4), diff --git a/api/services/rag_pipeline/rag_pipeline_transform_service.py b/api/services/rag_pipeline/rag_pipeline_transform_service.py index 39f426a2b0..d79ab71668 100644 --- a/api/services/rag_pipeline/rag_pipeline_transform_service.py +++ b/api/services/rag_pipeline/rag_pipeline_transform_service.py @@ -9,6 +9,7 @@ from flask_login import current_user from constants import DOCUMENT_EXTENSIONS from core.plugin.impl.plugin import PluginInstaller +from core.rag.retrieval.retrieval_methods import RetrievalMethod from extensions.ext_database import db from factories import variable_factory from models.dataset import Dataset, Document, DocumentPipelineExecutionLog, Pipeline @@ -164,7 +165,7 @@ class RagPipelineTransformService: if retrieval_model: retrieval_setting = RetrievalSetting.model_validate(retrieval_model) if indexing_technique == "economy": - retrieval_setting.search_method = "keyword_search" + retrieval_setting.search_method = RetrievalMethod.KEYWORD_SEARCH knowledge_configuration.retrieval_model = retrieval_setting else: dataset.retrieval_model = knowledge_configuration.retrieval_model.model_dump() diff --git a/api/tests/unit_tests/core/rag/extractor/firecrawl/test_firecrawl.py b/api/tests/unit_tests/core/rag/extractor/firecrawl/test_firecrawl.py index e5ead6ff66..b4ee1b91b4 100644 --- a/api/tests/unit_tests/core/rag/extractor/firecrawl/test_firecrawl.py +++ b/api/tests/unit_tests/core/rag/extractor/firecrawl/test_firecrawl.py @@ -1,10 +1,12 @@ import os +from pytest_mock import MockerFixture + from core.rag.extractor.firecrawl.firecrawl_app import FirecrawlApp from tests.unit_tests.core.rag.extractor.test_notion_extractor import _mock_response -def test_firecrawl_web_extractor_crawl_mode(mocker): +def test_firecrawl_web_extractor_crawl_mode(mocker: MockerFixture): url = "https://firecrawl.dev" api_key = os.getenv("FIRECRAWL_API_KEY") or "fc-" base_url = "https://api.firecrawl.dev" diff --git a/api/tests/unit_tests/core/rag/extractor/test_notion_extractor.py b/api/tests/unit_tests/core/rag/extractor/test_notion_extractor.py index f1e1820acc..58bec7d19e 100644 --- a/api/tests/unit_tests/core/rag/extractor/test_notion_extractor.py +++ b/api/tests/unit_tests/core/rag/extractor/test_notion_extractor.py @@ -1,5 +1,7 @@ from unittest import mock +from pytest_mock import MockerFixture + from core.rag.extractor import notion_extractor user_id = "user1" @@ -57,7 +59,7 @@ def _remove_multiple_new_lines(text): return text.strip() -def test_notion_page(mocker): +def test_notion_page(mocker: MockerFixture): texts = ["Head 1", "1.1", "paragraph 1", "1.1.1"] mocked_notion_page = { "object": "list", @@ -77,7 +79,7 @@ def test_notion_page(mocker): assert content == "# Head 1\n## 1.1\nparagraph 1\n### 1.1.1" -def test_notion_database(mocker): +def test_notion_database(mocker: MockerFixture): page_title_list = ["page1", "page2", "page3"] mocked_notion_database = { "object": "list", diff --git a/api/tests/unit_tests/core/test_model_manager.py b/api/tests/unit_tests/core/test_model_manager.py index d98e9f6bad..5a7547e85c 100644 --- a/api/tests/unit_tests/core/test_model_manager.py +++ b/api/tests/unit_tests/core/test_model_manager.py @@ -2,6 +2,7 @@ from unittest.mock import MagicMock, patch import pytest import redis +from pytest_mock import MockerFixture from core.entities.provider_entities import ModelLoadBalancingConfiguration from core.model_manager import LBModelManager @@ -39,7 +40,7 @@ def lb_model_manager(): return lb_model_manager -def test_lb_model_manager_fetch_next(mocker, lb_model_manager): +def test_lb_model_manager_fetch_next(mocker: MockerFixture, lb_model_manager: LBModelManager): # initialize redis client redis_client.initialize(redis.Redis()) diff --git a/api/tests/unit_tests/core/test_provider_manager.py b/api/tests/unit_tests/core/test_provider_manager.py index 2dab394029..0c3887beab 100644 --- a/api/tests/unit_tests/core/test_provider_manager.py +++ b/api/tests/unit_tests/core/test_provider_manager.py @@ -1,4 +1,5 @@ import pytest +from pytest_mock import MockerFixture from core.entities.provider_entities import ModelSettings from core.model_runtime.entities.model_entities import ModelType @@ -7,19 +8,25 @@ from models.provider import LoadBalancingModelConfig, ProviderModelSetting @pytest.fixture -def mock_provider_entity(mocker): +def mock_provider_entity(mocker: MockerFixture): mock_entity = mocker.Mock() mock_entity.provider = "openai" mock_entity.configurate_methods = ["predefined-model"] mock_entity.supported_model_types = [ModelType.LLM] - mock_entity.model_credential_schema = mocker.Mock() - mock_entity.model_credential_schema.credential_form_schemas = [] + # Use PropertyMock to ensure credential_form_schemas is iterable + provider_credential_schema = mocker.Mock() + type(provider_credential_schema).credential_form_schemas = mocker.PropertyMock(return_value=[]) + mock_entity.provider_credential_schema = provider_credential_schema + + model_credential_schema = mocker.Mock() + type(model_credential_schema).credential_form_schemas = mocker.PropertyMock(return_value=[]) + mock_entity.model_credential_schema = model_credential_schema return mock_entity -def test__to_model_settings(mocker, mock_provider_entity): +def test__to_model_settings(mocker: MockerFixture, mock_provider_entity): # Mocking the inputs provider_model_settings = [ ProviderModelSetting( @@ -79,7 +86,7 @@ def test__to_model_settings(mocker, mock_provider_entity): assert result[0].load_balancing_configs[1].name == "first" -def test__to_model_settings_only_one_lb(mocker, mock_provider_entity): +def test__to_model_settings_only_one_lb(mocker: MockerFixture, mock_provider_entity): # Mocking the inputs provider_model_settings = [ ProviderModelSetting( @@ -127,7 +134,7 @@ def test__to_model_settings_only_one_lb(mocker, mock_provider_entity): assert len(result[0].load_balancing_configs) == 0 -def test__to_model_settings_lb_disabled(mocker, mock_provider_entity): +def test__to_model_settings_lb_disabled(mocker: MockerFixture, mock_provider_entity): # Mocking the inputs provider_model_settings = [ ProviderModelSetting(