From e23b5f781ed41c75d27a6c7618cf29ccec6c14eb Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaan@berri.ai>
Date: Wed, 2 Sep 2026 20:56:12 -0700
Subject: [PATCH] test: add deterministic unit tests for vertex ocr url
 building and response wrapping

---
 tests/ocr_tests/test_ocr_vertex_ai.py | 71 ++++++++++++++++++++++++++-
 1 file changed, 70 insertions(+), 1 deletion(-)

diff --git a/tests/ocr_tests/test_ocr_vertex_ai.py b/tests/ocr_tests/test_ocr_vertex_ai.py
index 1842eb063a5..83dec39bc45 100644
--- a/tests/ocr_tests/test_ocr_vertex_ai.py
+++ b/tests/ocr_tests/test_ocr_vertex_ai.py
@@ -9,7 +9,9 @@ import json
 import os
 import tempfile
 from typing import Final
+from unittest.mock import MagicMock
 
+import httpx
 import pytest
 from base_ocr_unit_tests import BaseOCRTest
 
@@ -143,7 +145,9 @@ def test_vertex_ai_ocr_routing():
     ), "DeepSeek variant should route to VertexAIDeepSeekOCRConfig"
 
 
-@pytest.mark.parametrize("model", ("deepseek-ocr-maas", "deepseek-ai/deepseek-ocr-maas"))
+@pytest.mark.parametrize(
+    "model", ("deepseek-ocr-maas", "deepseek-ai/deepseek-ocr-maas")
+)
 def test_deepseek_request_uses_single_provider_namespace(model: str) -> None:
     from litellm.llms.vertex_ai.ocr.deepseek_transformation import (
         VertexAIDeepSeekOCRConfig,
@@ -157,3 +161,68 @@ def test_deepseek_request_uses_single_provider_namespace(model: str) -> None:
     )
 
     assert request.data["model"] == "deepseek-ai/deepseek-ocr-maas"
+
+
+def test_vertex_mistral_url_uses_project_location_and_model() -> None:
+    """Vertex Mistral OCR URL is built from project, location, and model in the Mistral publisher path."""
+    from litellm.llms.vertex_ai.ocr.transformation import VertexAIOCRConfig
+
+    url: Final = VertexAIOCRConfig().get_complete_url(
+        api_base=None,
+        model="mistral-ocr-maas",
+        optional_params={},
+        litellm_params={"vertex_project": "proj-1", "vertex_location": "europe-west4"},
+    )
+
+    assert url == (
+        "https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1"
+        "/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
+    )
+
+
+def test_vertex_mistral_reuses_mistral_body_transform() -> None:
+    """Vertex Mistral OCR delegates request-body construction to MistralOCRConfig unchanged."""
+    from litellm.llms.vertex_ai.ocr.transformation import VertexAIOCRConfig
+
+    request: Final = VertexAIOCRConfig().transform_ocr_request(
+        model="mistral-ocr-maas",
+        document={"type": "image_url", "image_url": "data:image/png;base64,abc"},
+        optional_params={},
+        headers={},
+    )
+
+    assert request.data["model"] == "mistral-ocr-maas"
+    assert request.data["document"]["image_url"] == "data:image/png;base64,abc"
+
+
+def test_vertex_deepseek_response_wraps_markdown_content() -> None:
+    """Vertex DeepSeek OCR wraps a plain-markdown chat completion into a single-page OCR response."""
+    from litellm.llms.vertex_ai.ocr.deepseek_transformation import (
+        VertexAIDeepSeekOCRConfig,
+    )
+
+    raw_response: Final = httpx.Response(
+        status_code=200,
+        json={
+            "choices": [{"message": {"content": "# OCR text"}}],
+            "usage": {"prompt_tokens": 1},
+        },
+    )
+
+    response = VertexAIDeepSeekOCRConfig().transform_ocr_response(
+        model="deepseek-ocr-maas",
+        raw_response=raw_response,
+        logging_obj=MagicMock(),
+    )
+
+    assert [page.model_dump() for page in response.pages] == [
+        {
+            "index": 0,
+            "markdown": "# OCR text",
+            "images": None,
+            "dimensions": None,
+        }
+    ]
+    assert response.model == "deepseek-ocr-maas"
+    assert response.usage_info is not None
+    assert response.usage_info.model_dump()["prompt_tokens"] == 1