diff --git a/api/tests/unit_tests/core/rag/rerank/test_reranker.py b/api/tests/unit_tests/core/rag/rerank/test_reranker.py index 7660da221c5..d138174a3c2 100644 --- a/api/tests/unit_tests/core/rag/rerank/test_reranker.py +++ b/api/tests/unit_tests/core/rag/rerank/test_reranker.py @@ -19,7 +19,9 @@ from unittest.mock import MagicMock, Mock, patch import pytest from sqlalchemy.orm import Session +from core.entities.provider_configuration import ProviderModelBundle from core.model_manager import ModelInstance +from core.plugin.impl.model_runtime_factory import create_plugin_model_runtime from core.rag.index_processor.constant.doc_type import DocType from core.rag.index_processor.constant.query_type import QueryType from core.rag.models.document import Document @@ -30,23 +32,35 @@ from core.rag.rerank.rerank_model import RerankModelRunner from core.rag.rerank.rerank_type import RerankMode from core.rag.rerank.weight_rerank import WeightRerankRunner from extensions.storage.storage_type import StorageType -from graphon.model_runtime.entities.model_entities import ModelFeature +from graphon.model_runtime.entities.model_entities import ModelFeature, ModelType from graphon.model_runtime.entities.rerank_entities import RerankDocument, RerankResult +from graphon.model_runtime.model_providers.base.rerank_model import RerankModel from models.enums import CreatorUserRole from models.model import UploadFile +from tests.unit_tests.core.model_fixtures import make_model_config -def create_mock_model_instance() -> ModelInstance: - """Create a properly configured mock ModelInstance for reranking tests.""" - mock_instance = Mock(spec=ModelInstance) - # Setup provider_model_bundle chain for check_model_support_vision - mock_instance.provider_model_bundle = Mock() - mock_instance.provider_model_bundle.configuration = Mock() - mock_instance.provider_model_bundle.configuration.tenant_id = "test-tenant-id" - mock_instance.provider = "test-provider" - mock_instance.model_name = "test-model" - mock_instance.get_model_schema.return_value = Mock(features=[]) - return mock_instance +def create_model_instance() -> ModelInstance: + """Build a rerank model with isolated schema and plugin invocation boundaries.""" + config = make_model_config(provider="test-provider", model="test-model", mode="chat") + configuration = config.provider_model_bundle.configuration + configuration.provider.supported_model_types = [ModelType.RERANK] + schema = config.model_schema.model_copy(update={"model_type": ModelType.RERANK, "features": []}) + instance = ModelInstance( + provider_model_bundle=ProviderModelBundle( + configuration=configuration, + model_type_instance=RerankModel( + provider_schema=configuration.provider, + model_runtime=create_plugin_model_runtime(tenant_id=configuration.tenant_id), + ), + ), + model=config.model, + credentials={}, + ) + instance.get_model_schema = Mock(return_value=schema) + instance.invoke_rerank = Mock() + instance.invoke_multimodal_rerank = Mock() + return instance class _UsesSQLiteSession: @@ -70,8 +84,8 @@ class TestRerankModelRunner(_UsesSQLiteSession): @pytest.fixture def mock_model_instance(self): - """Create a mock ModelInstance for reranking.""" - return create_mock_model_instance() + """Create the model with plugin I/O isolated.""" + return create_model_instance() @pytest.fixture def rerank_runner(self, mock_model_instance): @@ -428,7 +442,7 @@ class TestBaseRerankRunner: class TestRerankModelRunnerMultimodal(_UsesSQLiteSession): @pytest.fixture def mock_model_instance(self): - return create_mock_model_instance() + return create_model_instance() @pytest.fixture def rerank_runner(self, mock_model_instance): @@ -455,7 +469,8 @@ class TestRerankModelRunnerMultimodal(_UsesSQLiteSession): docs=[RerankDocument(index=0, text="doc", score=0.88)], ) - rerank_runner.rerank_model_instance.get_model_schema.return_value = Mock(features=[ModelFeature.VISION]) + schema = rerank_runner.rerank_model_instance.get_model_schema.return_value + schema.features = [ModelFeature.VISION] with patch.object( rerank_runner, "fetch_multimodal_rerank", @@ -1092,7 +1107,7 @@ class TestRerankRunnerFactory(_UsesSQLiteSession): - Parameters are forwarded to runner constructor """ # Arrange: Mock model instance - mock_model_instance = create_mock_model_instance() + mock_model_instance = create_model_instance() # Act: Create runner via factory runner = RerankRunnerFactory.create_rerank_runner( @@ -1155,7 +1170,7 @@ class TestRerankRunnerFactory(_UsesSQLiteSession): - String values are properly matched """ # Arrange: Mock model instance - mock_model_instance = create_mock_model_instance() + mock_model_instance = create_model_instance() # Act: Create runner using enum value runner = RerankRunnerFactory.create_rerank_runner( @@ -1186,7 +1201,7 @@ class TestRerankIntegration(_UsesSQLiteSession): - Top results are returned correctly """ # Arrange: Create mock model and documents - mock_model_instance = create_mock_model_instance() + mock_model_instance = create_model_instance() mock_rerank_result = RerankResult( model="bge-reranker-base", docs=[ @@ -1243,7 +1258,7 @@ class TestRerankIntegration(_UsesSQLiteSession): - Normalization is consistent """ # Arrange: Create mock model with various scores - mock_model_instance = create_mock_model_instance() + mock_model_instance = create_model_instance() mock_rerank_result = RerankResult( model="bge-reranker-base", docs=[ @@ -1292,7 +1307,7 @@ class TestRerankEdgeCases(_UsesSQLiteSession): - Empty metadata documents are processed correctly """ # Arrange: Create documents with empty metadata - mock_model_instance = create_mock_model_instance() + mock_model_instance = create_model_instance() mock_rerank_result = RerankResult( model="bge-reranker-base", docs=[ @@ -1338,7 +1353,7 @@ class TestRerankEdgeCases(_UsesSQLiteSession): - Score comparison logic works at boundary """ # Arrange: Create mock with various scores including negatives - mock_model_instance = create_mock_model_instance() + mock_model_instance = create_model_instance() mock_rerank_result = RerankResult( model="bge-reranker-base", docs=[ @@ -1374,7 +1389,7 @@ class TestRerankEdgeCases(_UsesSQLiteSession): - No overflow or precision issues """ # Arrange: All documents with perfect scores - mock_model_instance = create_mock_model_instance() + mock_model_instance = create_model_instance() mock_rerank_result = RerankResult( model="bge-reranker-base", docs=[ @@ -1409,7 +1424,7 @@ class TestRerankEdgeCases(_UsesSQLiteSession): - Content encoding is preserved """ # Arrange: Documents with special characters - mock_model_instance = create_mock_model_instance() + mock_model_instance = create_model_instance() mock_rerank_result = RerankResult( model="bge-reranker-base", docs=[ @@ -1451,7 +1466,7 @@ class TestRerankEdgeCases(_UsesSQLiteSession): - Content is not truncated unexpectedly """ # Arrange: Documents with very long content - mock_model_instance = create_mock_model_instance() + mock_model_instance = create_model_instance() long_content = "This is a very long document. " * 1000 # ~30,000 characters mock_rerank_result = RerankResult( @@ -1488,7 +1503,7 @@ class TestRerankEdgeCases(_UsesSQLiteSession): - All documents are processed correctly """ # Arrange: Create 100 documents - mock_model_instance = create_mock_model_instance() + mock_model_instance = create_model_instance() num_docs = 100 # Create rerank results for all documents @@ -1579,7 +1594,7 @@ class TestRerankEdgeCases(_UsesSQLiteSession): - Documents can still be ranked """ # Arrange: Empty query - mock_model_instance = create_mock_model_instance() + mock_model_instance = create_model_instance() mock_rerank_result = RerankResult( model="bge-reranker-base", docs=[ @@ -1626,7 +1641,7 @@ class TestRerankPerformance(_UsesSQLiteSession): - Efficient batch processing """ # Arrange: Multiple documents - mock_model_instance = create_mock_model_instance() + mock_model_instance = create_model_instance() mock_rerank_result = RerankResult( model="bge-reranker-base", docs=[RerankDocument(index=i, text=f"Doc {i}", score=0.9 - i * 0.1) for i in range(5)], @@ -1736,7 +1751,7 @@ class TestRerankErrorHandling(_UsesSQLiteSession): - Error context is preserved """ # Arrange: Mock model that raises exception - mock_model_instance = create_mock_model_instance() + mock_model_instance = create_model_instance() mock_model_instance.invoke_rerank.side_effect = RuntimeError("Model invocation failed") documents = [ @@ -1762,7 +1777,7 @@ class TestRerankErrorHandling(_UsesSQLiteSession): - Invalid results don't corrupt output """ # Arrange: Rerank result with invalid index - mock_model_instance = create_mock_model_instance() + mock_model_instance = create_model_instance() mock_rerank_result = RerankResult( model="bge-reranker-base", docs=[