From 18f75241e1e962334bac750cd91f412ad27fb677 Mon Sep 17 00:00:00 2001
From: Oguz Vuruskaner <ovuruska@outlook.com>
Date: Tue, 31 Mar 2026 12:24:08 +0300
Subject: [PATCH] test: update deepinfra rerank tests to expect single-element
 queries array

Update 4 test assertions across 3 test files to match the new
single-element queries behavior. All 33 tests pass.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .../llms/deepinfra/test_deepinfra_rerank.py   |  3 +--
 .../test_deepinfra_rerank_integration.py      |  3 +--
 .../test_deepinfra_rerank_transformation.py   | 26 ++++++-------------
 3 files changed, 10 insertions(+), 22 deletions(-)

diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank.py b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank.py
index 0dda7d08da4..ec47f3a2dc4 100644
--- a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank.py
+++ b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank.py
@@ -251,8 +251,7 @@ def test_deepinfra_rerank_request_format(mock_post):
     request_data = json.loads(mock_post.call_args.kwargs["data"])
     assert request_data["queries"] == [
         "test query",
-        "test query",
-    ]  # DeepInfra requires queries to match documents length
+    ]  # Single-element array; Deepinfra broadcasts across all documents
     assert request_data["documents"] == ["doc1", "doc2"]
     assert request_data["instruction"] == "custom instruction"
     assert request_data["webhook"] == "https://webhook.example.com"
diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_integration.py b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_integration.py
index 3655f5c643b..a428cac24ec 100644
--- a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_integration.py
+++ b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_integration.py
@@ -369,8 +369,7 @@ def test_deepinfra_rerank_request_format(mock_post):
     request_data = json.loads(mock_post.call_args.kwargs["data"])
     assert request_data["queries"] == [
         "test query",
-        "test query",
-    ]  # DeepInfra requires queries to match documents length
+    ]  # Single-element array; Deepinfra broadcasts across all documents
     assert request_data["documents"] == ["doc1", "doc2"]
     assert request_data["instruction"] == "custom instruction"
     assert request_data["webhook"] == "https://webhook.example.com"
diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py
index 252eb40532c..c69ea599c47 100644
--- a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py
+++ b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py
@@ -54,8 +54,7 @@ class TestDeepinfraRerankTransform:
         )
         assert params["queries"] == [
             "test query",
-            "test query",
-        ]  # DeepInfra requires queries to match documents length
+        ]  # Single-element array; Deepinfra broadcasts across all documents
         assert params["documents"] == ["doc1", "doc2"]
 
     def test_map_cohere_rerank_params_with_non_default(self):
@@ -228,16 +227,10 @@ class TestDeepinfraRerankTransform:
         assert "documents" in supported_params
         assert len(supported_params) == 2
 
-    def test_query_replication_for_deepinfra_requirement(self):
-        """Test that queries are replicated to match documents length as required by DeepInfra."""
-        # Test with different document lengths
-        test_cases = [
-            (["doc1"], ["query1"]),
-            (["doc1", "doc2"], ["query1", "query1"]),
-            (["doc1", "doc2", "doc3"], ["query1", "query1", "query1"]),
-        ]
-
-        for documents, expected_queries in test_cases:
+    def test_query_single_element_array_for_deepinfra(self):
+        """Test that queries is a single-element array regardless of document count; Deepinfra broadcasts it."""
+        for num_docs in [1, 2, 3, 10]:
+            documents = [f"doc{i}" for i in range(num_docs)]
             params = self.config.map_cohere_rerank_params(
                 non_default_params={},
                 model=self.model,
@@ -245,12 +238,9 @@ class TestDeepinfraRerankTransform:
                 query="query1",
                 documents=documents,
             )
-            assert (
-                params["queries"] == expected_queries
-            ), f"Failed for {len(documents)} documents"
-            assert len(params["queries"]) == len(
-                documents
-            ), "Queries length must match documents length"
+            assert params["queries"] == [
+                "query1"
+            ], f"Failed for {num_docs} documents"
 
     def test_get_error_class_basic(self):
         """Test error class generation for basic error."""