diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index 64ae81b5763..144ba2b7b6d 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/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,")