test(batches): move orphan tests into tests/test_litellm for CI coverage

tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py → tests/test_litellm/llms/vertex_ai/test_gemini_batch_embeddings.py

@@ -8,12 +8,8 @@
 """
 
 import json
-import os
-import sys
 from unittest.mock import MagicMock, patch
 
-sys.path.insert(0, os.path.abspath("../../../.."))
-
 import pytest
 
 import litellm
@@ -311,13 +307,16 @@ def mock_auth_token(*args, **kwargs):
     ):
         mock_get_token.return_value = (
             {"x-goog-api-key": "test-key"},
-            "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:embedContent",
+            "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:batchEmbedContents",
         )
 
         mock_response = MagicMock()
         mock_response.status_code = 200
         mock_response.json.return_value = {
-            "embedding": {"values": [0.1, 0.2, 0.3, 0.4, 0.5]}
+            "embeddings": [
+                {"values": [0.1, 0.2, 0.3, 0.4, 0.5]},
+                {"values": [0.6, 0.7, 0.8, 0.9, 1.0]},
+            ]
         }
         mock_post.return_value = mock_response
 
@@ -338,17 +337,21 @@ def mock_auth_token(*args, **kwargs):
 
         request_body = json.loads(kwargs.get("data", "{}"))
 
-        assert "content" in request_body
-        assert "parts" in request_body["content"]
-        parts = request_body["content"]["parts"]
+        assert "requests" in request_body
+        assert len(request_body["requests"]) == 2
 
-        assert len(parts) == 2
-        assert parts[0]["text"] == "The food was delicious"
-        assert "inline_data" in parts[1]
-        assert parts[1]["inline_data"]["mime_type"] == "image/png"
+        text_parts = request_body["requests"][0]["content"]["parts"]
+        image_parts = request_body["requests"][1]["content"]["parts"]
 
-        assert len(response.data) == 1
+        assert len(text_parts) == 1
+        assert text_parts[0]["text"] == "The food was delicious"
+        assert len(image_parts) == 1
+        assert "inline_data" in image_parts[0]
+        assert image_parts[0]["inline_data"]["mime_type"] == "image/png"
+
+        assert len(response.data) == 2
         assert response.data[0].embedding == [0.1, 0.2, 0.3, 0.4, 0.5]
+        assert response.data[1].embedding == [0.6, 0.7, 0.8, 0.9, 1.0]
 
 
 def test_gemini_multimodal_embedding_with_audio():
@@ -581,17 +584,21 @@ def mock_auth_token(*args, **kwargs):
 
 def test_filter_embed_params_drops_unsupported():
     """Unsupported params like max_tokens should be filtered out."""
-    result = _filter_embed_params({"dimensions": 768, "max_tokens": 256, "temperature": 0.5})
+    result = _filter_embed_params(
+        {"dimensions": 768, "max_tokens": 256, "temperature": 0.5}
+    )
     assert result == {"outputDimensionality": 768}
 
 
 def test_filter_embed_params_keeps_supported():
     """All supported Gemini embedding params should pass through."""
-    result = _filter_embed_params({
-        "dimensions": 768,
-        "task_type": "RETRIEVAL_DOCUMENT",
-        "title": "My doc",
-    })
+    result = _filter_embed_params(
+        {
+            "dimensions": 768,
+            "task_type": "RETRIEVAL_DOCUMENT",
+            "title": "My doc",
+        }
+    )
     assert result == {
         "outputDimensionality": 768,
         "taskType": "RETRIEVAL_DOCUMENT",

tests/litellm/proxy/test_batch_x_litellm_model_encoding.py → tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py

@@ -422,9 +422,7 @@ async def test_cancel_batch_with_unified_id_routes_with_decoded_model_and_batch_
 
     model_id = "deployment-123"
     raw_batch_id = "batch_openai_123"
-    unified_batch_id = _make_unified_batch_id(
-        model_id=model_id, batch_id=raw_batch_id
-    )
+    unified_batch_id = _make_unified_batch_id(model_id=model_id, batch_id=raw_batch_id)
     mock_response = _make_batch_response(batch_id=raw_batch_id, status="cancelled")
     mock_response._hidden_params = {}
     mock_router = MagicMock()