test_image_exceeds_size_limit_with_content_length

tests/test_litellm/litellm_core_utils/test_image_handling.py

@@ -1,7 +1,10 @@
+from unittest.mock import patch
+
 import pytest
 from httpx import Request, Response
 
 import litellm
+from litellm import constants
 from litellm.litellm_core_utils.prompt_templates.image_handling import (
     convert_url_to_base64,
 )
@@ -39,3 +42,96 @@ def test_completion_with_invalid_image_url(monkeypatch):
         )
     assert excinfo.value.status_code == 400
     assert "Unable to fetch image" in str(excinfo.value)
+
+
+class LargeImageClient:
+    """
+    Client that returns a large image exceeding size limit.
+    """
+
+    def __init__(self, size_mb=100, include_content_length=True):
+        self.size_mb = size_mb
+        self.include_content_length = include_content_length
+
+    def get(self, url, follow_redirects=True):
+        size_bytes = int(self.size_mb * 1024 * 1024)
+        headers = {"Content-Type": "image/jpeg"}
+        if self.include_content_length:
+            headers["Content-Length"] = str(size_bytes)
+        return Response(
+            status_code=200,
+            headers=headers,
+            content=b"x" * size_bytes,
+            request=Request("GET", url),
+        )
+
+
+def test_image_exceeds_size_limit_with_content_length(monkeypatch):
+    """
+    Test that images exceeding MAX_IMAGE_URL_DOWNLOAD_SIZE_MB are rejected when Content-Length header is present.
+    """
+    monkeypatch.setattr(litellm, "module_level_client", LargeImageClient(size_mb=100))
+
+    with pytest.raises(litellm.ImageFetchError) as excinfo:
+        convert_url_to_base64("https://example.com/large-image.jpg")
+
+    assert "exceeds maximum allowed size" in str(excinfo.value)
+    assert "100.00MB" in str(excinfo.value)
+    assert "50.0MB" in str(excinfo.value)
+
+
+def test_image_exceeds_size_limit_without_content_length(monkeypatch):
+    """
+    Test that images exceeding MAX_IMAGE_URL_DOWNLOAD_SIZE_MB are rejected even without Content-Length header.
+    """
+    monkeypatch.setattr(
+        litellm, "module_level_client", LargeImageClient(size_mb=100, include_content_length=False)
+    )
+
+    with pytest.raises(litellm.ImageFetchError) as excinfo:
+        convert_url_to_base64("https://example.com/large-image.jpg")
+
+    assert "exceeds maximum allowed size" in str(excinfo.value)
+
+
+class SmallImageClient:
+    """
+    Client that returns a small valid image.
+    """
+
+    def get(self, url, follow_redirects=True):
+        size_bytes = 1024
+        headers = {
+            "Content-Type": "image/jpeg",
+            "Content-Length": str(size_bytes),
+        }
+        return Response(
+            status_code=200,
+            headers=headers,
+            content=b"x" * size_bytes,
+            request=Request("GET", url),
+        )
+
+
+def test_image_within_size_limit(monkeypatch):
+    """
+    Test that images within size limit are processed successfully.
+    """
+    monkeypatch.setattr(litellm, "module_level_client", SmallImageClient())
+
+    result = convert_url_to_base64("https://example.com/small-image.jpg")
+
+    assert result.startswith("data:image/jpeg;base64,")
+
+
+def test_image_size_limit_disabled(monkeypatch):
+    """
+    Test that setting MAX_IMAGE_URL_DOWNLOAD_SIZE_MB to 0 disables size checking.
+    """
+    import litellm.litellm_core_utils.prompt_templates.image_handling as image_handling
+
+    monkeypatch.setattr(litellm, "module_level_client", LargeImageClient(size_mb=100))
+    monkeypatch.setattr(image_handling, "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB", 0)
+
+    result = convert_url_to_base64("https://example.com/large-image.jpg")
+    assert result.startswith("data:image/jpeg;base64,")