mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-07 04:24:12 +00:00
Merge pull request #21268 from BerriAI/fix/vertex-ai-rerank-test-isolation
fix(test): add environment cleanup for Vertex AI rerank tests
This commit is contained in:
@@ -15,9 +15,30 @@ from litellm.types.rerank import RerankResponse
|
||||
|
||||
class TestVertexAIRerankTransform:
|
||||
def setup_method(self):
|
||||
# Save and clear Google/Vertex AI environment variables to prevent
|
||||
# test isolation issues where previous tests leave credentials set
|
||||
self._saved_env = {}
|
||||
env_vars_to_clear = [
|
||||
"GOOGLE_APPLICATION_CREDENTIALS",
|
||||
"GOOGLE_CLOUD_PROJECT",
|
||||
"VERTEXAI_PROJECT",
|
||||
"VERTEX_PROJECT",
|
||||
"VERTEX_LOCATION",
|
||||
"VERTEX_AI_PROJECT",
|
||||
]
|
||||
for var in env_vars_to_clear:
|
||||
if var in os.environ:
|
||||
self._saved_env[var] = os.environ[var]
|
||||
del os.environ[var]
|
||||
|
||||
self.config = VertexAIRerankConfig()
|
||||
self.model = "semantic-ranker-default@latest"
|
||||
|
||||
def teardown_method(self):
|
||||
# Restore saved environment variables
|
||||
for var, value in self._saved_env.items():
|
||||
os.environ[var] = value
|
||||
|
||||
@patch('litellm.llms.vertex_ai.rerank.transformation.VertexAIRerankConfig._ensure_access_token')
|
||||
def test_get_complete_url(self, mock_ensure_access_token):
|
||||
"""Test URL generation for Vertex AI Discovery Engine rerank API."""
|
||||
|
||||
Reference in New Issue
Block a user