diff --git a/tests/llm_translation/test_prompt_factory.py b/tests/llm_translation/test_prompt_factory.py index 36e47e3c2f4..05a58a135d2 100644 --- a/tests/llm_translation/test_prompt_factory.py +++ b/tests/llm_translation/test_prompt_factory.py @@ -605,8 +605,33 @@ def test_no_messages_yields_user_text(): assert contents == expected_output -def test_convert_url(): - convert_url_to_base64("https://picsum.photos/id/237/200/300") +def test_convert_url(monkeypatch): + import base64 + from unittest.mock import MagicMock + + import httpx + + from litellm.litellm_core_utils.prompt_templates.image_handling import ( + in_memory_cache, + ) + + url = "https://picsum.photos/id/237/200/300" + image_bytes = b"\x89PNG\r\n\x1a\nfake-png-bytes" + + mock_client = MagicMock() + mock_client.get.return_value = httpx.Response( + 200, content=image_bytes, headers={"Content-Type": "image/png"} + ) + + monkeypatch.setattr(litellm, "user_url_validation", False, raising=False) + monkeypatch.setattr(litellm, "module_level_client", mock_client, raising=False) + in_memory_cache.flush_cache() + + result = convert_url_to_base64(url) + + expected = "data:image/png;base64," + base64.b64encode(image_bytes).decode("utf-8") + assert result == expected + mock_client.get.assert_called_once() def test_azure_tool_call_invoke_helper():