Fix Vertex AI rerank tests for parallel test execution

The test_end_to_end_rerank_flow mock for _ensure_access_token was not
being applied because conftest reloads litellm, causing the
VertexAIRerankConfig class to be a different object than what's patched.

Fix: Reload the transformation module in setup_method and re-import the
class to ensure the patch targets the same class object used by tests.
This commit is contained in:
Julio Quinteros Pro 2026-01-26 03:07:26 -03:00
parent 8285a2a7b4
commit bee4c94556

View file

@ -2,6 +2,7 @@
Integration tests for Vertex AI rerank functionality.
These tests demonstrate end-to-end usage of the Vertex AI rerank feature.
"""
import importlib
import os
from unittest.mock import MagicMock, patch
@ -13,7 +14,14 @@ from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig
class TestVertexAIRerankIntegration:
def setup_method(self):
self.config = VertexAIRerankConfig()
# Reload modules to ensure fresh references after conftest reloads litellm.
# This ensures the class being patched is the same one used by the tests.
import litellm.llms.vertex_ai.rerank.transformation as rerank_transformation_module
importlib.reload(rerank_transformation_module)
# Re-import after reload to get the fresh class
from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig as FreshConfig
self.config = FreshConfig()
self.model = "semantic-ranker-default@latest"
@patch('litellm.llms.vertex_ai.rerank.transformation.VertexAIRerankConfig._ensure_access_token')