| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import asyncio |
| 5 | import mimetypes |
| 6 | import os |
| 7 | import shutil |
| 8 | import time |
| 9 | from io import BytesIO |
| 10 | from tempfile import NamedTemporaryFile, TemporaryDirectory |
| 11 | |
| 12 | import aiohttp |
| 13 | import numpy as np |
| 14 | import pybase64 as base64 |
| 15 | import pytest |
| 16 | import requests |
| 17 | import torch |
| 18 | from PIL import Image, ImageChops |
| 19 | |
| 20 | from vllm.assets.base import VLLM_S3_BUCKET_URL |
| 21 | from vllm.multimodal.image import convert_image_mode |
| 22 | from vllm.multimodal.inputs import PlaceholderRange |
| 23 | from vllm.multimodal.media import MediaConnector |
| 24 | |
| 25 | # Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA) |
| 26 | TEST_IMAGE_ASSETS = [ |
| 27 | "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg" |
| 28 | "Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png", |
| 29 | "1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png", |
| 30 | "RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png", |
| 31 | ] |
| 32 | |
| 33 | TEST_VIDEO_URLS = [ |
| 34 | f"{VLLM_S3_BUCKET_URL}/multimodal_asset/slow_traffic_small.mp4", |
| 35 | f"{VLLM_S3_BUCKET_URL}/multimodal_asset/vtest.avi", |
| 36 | ] |
| 37 | |
| 38 | |
| 39 | @pytest.fixture(scope="module") |
| 40 | def url_images(local_asset_server) -> dict[str, Image.Image]: |
| 41 | return { |
| 42 | image_url: local_asset_server.get_image_asset(image_url) |
| 43 | for image_url in TEST_IMAGE_ASSETS |
| 44 | } |
| 45 | |
| 46 | |
| 47 | def get_supported_suffixes() -> tuple[str, ...]: |
| 48 | # We should at least test the file types mentioned in GPT-4 with Vision |
| 49 | OPENAI_SUPPORTED_SUFFIXES = (".png", ".jpeg", ".jpg", ".webp", ".gif") |
| 50 | |
| 51 | # Additional file types that are supported by us |
| 52 | EXTRA_SUPPORTED_SUFFIXES = (".bmp", ".tiff") |
| 53 | |
| 54 | return OPENAI_SUPPORTED_SUFFIXES + EXTRA_SUPPORTED_SUFFIXES |
| 55 | |
| 56 | |
| 57 | def _image_equals(a: Image.Image, b: Image.Image) -> bool: |
| 58 | return (np.asarray(a) == np.asarray(convert_image_mode(b, a.mode))).all() |
| 59 | |
| 60 | |
| 61 | @pytest.mark.asyncio |
| 62 | @pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True) |
| 63 | async def test_fetch_image_http(image_url: str): |
| 64 | connector = MediaConnector() |
| 65 | |
| 66 | image_sync = connector.fetch_image(image_url) |
| 67 | image_async = await connector.fetch_image_async(image_url) |
| 68 | assert _image_equals(image_sync, image_async) |
| 69 | |
| 70 | |
| 71 | @pytest.mark.asyncio |
| 72 | @pytest.mark.parametrize("raw_image_url", TEST_IMAGE_ASSETS) |
| 73 | @pytest.mark.parametrize("suffix", get_supported_suffixes()) |
| 74 | async def test_fetch_image_base64( |
| 75 | url_images: dict[str, Image.Image], raw_image_url: str, suffix: str |
| 76 | ): |
| 77 | connector = MediaConnector( |
| 78 | # Domain restriction should not apply to data URLs. |
| 79 | allowed_media_domains=[ |
| 80 | VLLM_S3_BUCKET_URL.removeprefix("https://"), |
| 81 | ] |
| 82 | ) |
| 83 | url_image = url_images[raw_image_url] |
| 84 | |
| 85 | try: |
| 86 | mime_type = Image.MIME[Image.registered_extensions()[suffix]] |
| 87 | except KeyError: |
| 88 | try: |
| 89 | mime_type = mimetypes.types_map[suffix] |
| 90 | except KeyError: |
| 91 | pytest.skip("No MIME type") |
| 92 | |
| 93 | with NamedTemporaryFile(suffix=suffix) as f: |
| 94 | try: |
| 95 | url_image.save(f.name) |
| 96 | except Exception as e: |
| 97 | if e.args[0] == "cannot write mode RGBA as JPEG": |
| 98 | pytest.skip("Conversion not supported") |
| 99 | |
| 100 | raise |
| 101 | |
| 102 | base64_image = base64.b64encode(f.read()).decode("utf-8") |
| 103 | data_url = f"data:{mime_type};base64,{base64_image}" |
| 104 | |
| 105 | data_image_sync = connector.fetch_image(data_url) |
| 106 | if _image_equals(url_image, Image.open(f)): |
| 107 | assert _image_equals(url_image, data_image_sync) |
| 108 | else: |
| 109 | pass # Lossy format; only check that image can be opened |
| 110 | |
| 111 | data_image_async = await connector.fetch_image_async(data_url) |
| 112 | assert _image_equals(data_image_sync, data_image_async) |
| 113 | |
| 114 | |
| 115 | @pytest.mark.asyncio |
| 116 | async def test_fetch_image_keep_original_mode(): |
| 117 | """media_io_kwargs can disable the default RGB conversion.""" |
| 118 | # RGBA image: opaque black pixel on a fully transparent background |
| 119 | rgba_image = Image.new("RGBA", (4, 4), (0, 0, 0, 0)) |
| 120 | rgba_image.putpixel((2, 2), (0, 0, 0, 255)) |
| 121 | buffer = BytesIO() |
| 122 | rgba_image.save(buffer, "PNG") |
| 123 | data_url = ( |
| 124 | f"data:image/png;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}" |
| 125 | ) |
| 126 | |
| 127 | # Default behavior: RGBA is composited onto a white background |
| 128 | default_image = MediaConnector().fetch_image(data_url) |
| 129 | assert default_image.mode == "RGB" |
| 130 | assert default_image.getpixel((0, 0)) == (255, 255, 255) |
| 131 | assert default_image.getpixel((2, 2)) == (0, 0, 0) |
| 132 | |
| 133 | # image_mode=None via media_io_kwargs: original mode is preserved |
| 134 | connector = MediaConnector(media_io_kwargs={"image": {"image_mode": None}}) |
| 135 | image_sync = connector.fetch_image(data_url) |
| 136 | image_async = await connector.fetch_image_async(data_url) |
| 137 | for image in (image_sync, image_async): |
| 138 | assert image.mode == "RGBA" |
| 139 | assert image.getpixel((0, 0)) == (0, 0, 0, 0) |
| 140 | assert image.getpixel((2, 2)) == (0, 0, 0, 255) |
| 141 | |
| 142 | |
| 143 | @pytest.mark.asyncio |
| 144 | @pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True) |
| 145 | async def test_fetch_image_local_files(image_url: str): |
| 146 | connector = MediaConnector() |
| 147 | |
| 148 | with TemporaryDirectory() as temp_dir: |
| 149 | local_connector = MediaConnector(allowed_local_media_path=temp_dir) |
| 150 | |
| 151 | origin_image = connector.fetch_image(image_url) |
| 152 | origin_image.save( |
| 153 | os.path.join(temp_dir, os.path.basename(image_url)), |
| 154 | quality=100, |
| 155 | icc_profile=origin_image.info.get("icc_profile"), |
| 156 | ) |
| 157 | |
| 158 | image_async = await local_connector.fetch_image_async( |
| 159 | f"file://{temp_dir}/{os.path.basename(image_url)}" |
| 160 | ) |
| 161 | image_sync = local_connector.fetch_image( |
| 162 | f"file://{temp_dir}/{os.path.basename(image_url)}" |
| 163 | ) |
| 164 | # Check that the images are equal |
| 165 | assert not ImageChops.difference(image_sync, image_async).getbbox() |
| 166 | |
| 167 | with pytest.raises(ValueError, match="must be a subpath"): |
| 168 | await local_connector.fetch_image_async( |
| 169 | f"file://{temp_dir}/../{os.path.basename(image_url)}" |
| 170 | ) |
| 171 | with pytest.raises(RuntimeError, match="Cannot load local files"): |
| 172 | await connector.fetch_image_async( |
| 173 | f"file://{temp_dir}/../{os.path.basename(image_url)}" |
| 174 | ) |
| 175 | |
| 176 | with pytest.raises(ValueError, match="must be a subpath"): |
| 177 | local_connector.fetch_image( |
| 178 | f"file://{temp_dir}/../{os.path.basename(image_url)}" |
| 179 | ) |
| 180 | with pytest.raises(RuntimeError, match="Cannot load local files"): |
| 181 | connector.fetch_image(f"file://{temp_dir}/../{os.path.basename(image_url)}") |
| 182 | |
| 183 | |
| 184 | @pytest.mark.asyncio |
| 185 | async def test_fetch_image_local_files_relative_allowed_path(tmp_path, monkeypatch): |
| 186 | media_dir = tmp_path / "media" |
| 187 | media_dir.mkdir() |
| 188 | image_path = media_dir / "image.png" |
| 189 | Image.new("RGB", (1, 1), color=(255, 0, 0)).save(image_path) |
| 190 | |
| 191 | monkeypatch.chdir(tmp_path) |
| 192 | local_connector = MediaConnector(allowed_local_media_path="media") |
| 193 | |
| 194 | image_sync = local_connector.fetch_image(image_path.as_uri()) |
| 195 | image_async = await local_connector.fetch_image_async(image_path.as_uri()) |
| 196 | |
| 197 | assert image_sync.size == (1, 1) |
| 198 | assert not ImageChops.difference(image_sync, image_async).getbbox() |
| 199 | |
| 200 | |
| 201 | @pytest.mark.asyncio |
| 202 | @pytest.mark.parametrize("image_url", [TEST_IMAGE_ASSETS[0]], indirect=True) |
| 203 | async def test_fetch_image_local_files_with_space_in_name(image_url: str): |
| 204 | connector = MediaConnector() |
| 205 | |
| 206 | with TemporaryDirectory() as temp_dir: |
| 207 | local_connector = MediaConnector(allowed_local_media_path=temp_dir) |
| 208 | |
| 209 | origin_image = connector.fetch_image(image_url) |
| 210 | filename = "file name with space.jpg" |
| 211 | origin_image.save( |
| 212 | os.path.join(temp_dir, filename), |
| 213 | quality=100, |
| 214 | icc_profile=origin_image.info.get("icc_profile"), |
| 215 | ) |
| 216 | |
| 217 | try: |
| 218 | image_async = await local_connector.fetch_image_async( |
| 219 | f"file://{temp_dir}/{filename}" |
| 220 | ) |
| 221 | image_sync = local_connector.fetch_image(f"file://{temp_dir}/{filename}") |
| 222 | except FileNotFoundError as e: |
| 223 | pytest.fail("Failed to fetch image with space in name: {}".format(e)) |
| 224 | # Check that the images are equal |
| 225 | assert not ImageChops.difference(image_sync, image_async).getbbox() |
| 226 | |
| 227 | |
| 228 | @pytest.mark.asyncio |
| 229 | async def test_fetch_image_data_url_with_params(): |
| 230 | """RFC 2397 allows parameters between the mediatype and the base64 |
| 231 | marker; they must not be rejected or leak into the media type.""" |
| 232 | connector = MediaConnector() |
| 233 | |
| 234 | image = Image.new("RGB", (4, 4), color=(255, 0, 0)) |
| 235 | with NamedTemporaryFile(suffix=".png") as f: |
| 236 | image.save(f.name) |
| 237 | base64_image = base64.b64encode(f.read()).decode("utf-8") |
| 238 | |
| 239 | data_url = f"data:image/png;charset=utf-8;base64,{base64_image}" |
| 240 | image_sync = connector.fetch_image(data_url) |
| 241 | image_async = await connector.fetch_image_async(data_url) |
| 242 | assert _image_equals(image_sync, image_async) |
| 243 | |
| 244 | |
| 245 | def test_fetch_image_data_url_malformed(): |
| 246 | connector = MediaConnector() |
| 247 | |
| 248 | with pytest.raises(ValueError, match="missing ','"): |
| 249 | connector.fetch_image("data:image/png;base64") |
| 250 | |
| 251 | with pytest.raises(NotImplementedError, match="base64"): |
| 252 | connector.fetch_image("data:text/plain,hello") |
| 253 | |
| 254 | # ";base64" requires the ";"; here "base64" is a (bogus) media type. |
| 255 | with pytest.raises(NotImplementedError, match="base64"): |
| 256 | connector.fetch_image("data:base64,aGVsbG8=") |
| 257 | |
| 258 | # Strict RFC 2397 grammar: lowercase "base64", no whitespace. |
| 259 | with pytest.raises(NotImplementedError, match="base64"): |
| 260 | connector.fetch_image("data:image/png;BASE64,aGVsbG8=") |
| 261 | |
| 262 | with pytest.raises(NotImplementedError, match="base64"): |
| 263 | connector.fetch_image("data:image/png; base64,aGVsbG8=") |
| 264 | |
| 265 | |
| 266 | @pytest.mark.asyncio |
| 267 | async def test_fetch_image_error_conversion(): |
| 268 | connector = MediaConnector() |
| 269 | broken_img = "data:image/png;base64,aGVsbG9fdmxsbV9jb21tdW5pdHkK" |
| 270 | |
| 271 | # PIL.UnidentifiedImageError should be converted to ValueError |
| 272 | with pytest.raises(ValueError): |
| 273 | await connector.fetch_image_async(broken_img) |
| 274 | |
| 275 | with pytest.raises(ValueError): |
| 276 | connector.fetch_image(broken_img) |
| 277 | |
| 278 | |
| 279 | @pytest.mark.flaky(reruns=3, reruns_delay=5) |
| 280 | @pytest.mark.asyncio |
| 281 | @pytest.mark.parametrize("video_url", TEST_VIDEO_URLS) |
| 282 | @pytest.mark.parametrize("num_frames", [-1, 32, 1800]) |
| 283 | async def test_fetch_video_http(video_url: str, num_frames: int): |
| 284 | connector = MediaConnector( |
| 285 | media_io_kwargs={ |
| 286 | "video": { |
| 287 | "num_frames": num_frames, |
| 288 | } |
| 289 | } |
| 290 | ) |
| 291 | |
| 292 | try: |
| 293 | video_sync, metadata_sync = connector.fetch_video(video_url) |
| 294 | video_async, metadata_async = await connector.fetch_video_async(video_url) |
| 295 | except (TimeoutError, asyncio.TimeoutError) as e: |
| 296 | pytest.skip(f"Timeout fetching video (CI network flakiness): {e}") |
| 297 | |
| 298 | assert np.array_equal(video_sync, video_async) |
| 299 | assert metadata_sync == metadata_async |
| 300 | |
| 301 | |
| 302 | @pytest.mark.flaky(reruns=3, reruns_delay=5) |
| 303 | @pytest.mark.asyncio |
| 304 | @pytest.mark.parametrize("video_url", TEST_VIDEO_URLS) |
| 305 | @pytest.mark.parametrize("max_duration", [1, 60, 1800]) |
| 306 | @pytest.mark.parametrize("requested_fps", [2, 24]) |
| 307 | async def test_fetch_video_http_with_dynamic_loader( |
| 308 | video_url: str, |
| 309 | max_duration: int, |
| 310 | requested_fps: int, |
| 311 | monkeypatch: pytest.MonkeyPatch, |
| 312 | ): |
| 313 | with monkeypatch.context() as m: |
| 314 | m.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv_dynamic") |
| 315 | connector = MediaConnector( |
| 316 | media_io_kwargs={ |
| 317 | "video": { |
| 318 | "max_duration": max_duration, |
| 319 | "requested_fps": requested_fps, |
| 320 | } |
| 321 | } |
| 322 | ) |
| 323 | |
| 324 | try: |
| 325 | video_sync, metadata_sync = connector.fetch_video(video_url) |
| 326 | video_async, metadata_async = await connector.fetch_video_async(video_url) |
| 327 | except (TimeoutError, asyncio.TimeoutError) as e: |
| 328 | pytest.skip(f"Timeout fetching video (CI network flakiness): {e}") |
| 329 | |
| 330 | assert np.array_equal(video_sync, video_async) |
| 331 | assert metadata_sync == metadata_async |
| 332 | assert metadata_sync["video_backend"] == "opencv_dynamic" |
| 333 | |
| 334 | |
| 335 | @pytest.mark.parametrize( |
| 336 | "is_embed,start_idx,end_idx,expected", |
| 337 | [ |
| 338 | (None, 2, 4, (2, 4)), |
| 339 | ( |
| 340 | torch.tensor([False, True, False, True, True]), |
| 341 | 3, |
| 342 | 5, |
| 343 | (1, 3), |
| 344 | ), |
| 345 | ( |
| 346 | torch.tensor([False, True, False, True, True]), |
| 347 | 0, |
| 348 | 2, |
| 349 | (0, 1), |
| 350 | ), |
| 351 | ( |
| 352 | torch.tensor([True, False, True, False]), |
| 353 | 2, |
| 354 | 2, |
| 355 | (1, 1), |
| 356 | ), |
| 357 | ], |
| 358 | ) |
| 359 | def test_placeholder_range_get_embeds_indices_in_range( |
| 360 | is_embed, start_idx, end_idx, expected |
| 361 | ): |
| 362 | length = len(is_embed) if is_embed is not None else 5 |
| 363 | pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed) |
| 364 | assert pr.get_embeds_indices_in_range(start_idx, end_idx) == expected |
| 365 | |
| 366 | |
| 367 | @pytest.mark.parametrize( |
| 368 | "offset,is_embed,expected", |
| 369 | [ |
| 370 | (0, None, [(0, 4)]), |
| 371 | ( |
| 372 | 2, |
| 373 | torch.tensor([False, True, False, True, True]), |
| 374 | [(3, 3), (5, 6)], |
| 375 | ), |
| 376 | (0, torch.tensor([True, True, True, True]), [(0, 3)]), |
| 377 | (0, torch.tensor([False, False, False, False]), []), |
| 378 | ], |
| 379 | ) |
| 380 | def test_placeholder_range_extract_embeds_range(offset, is_embed, expected): |
| 381 | length = len(is_embed) if is_embed is not None else 5 |
| 382 | pr = PlaceholderRange(offset=offset, length=length, is_embed=is_embed) |
| 383 | assert pr.extract_embeds_range() == expected |
| 384 | |
| 385 | |
| 386 | @pytest.mark.flaky(reruns=3, reruns_delay=5) |
| 387 | @pytest.mark.asyncio |
| 388 | @pytest.mark.parametrize("video_url", TEST_VIDEO_URLS) |
| 389 | @pytest.mark.parametrize("num_frames", [-1, 32, 1800]) |
| 390 | async def test_allowed_media_domains(video_url: str, num_frames: int): |
| 391 | connector = MediaConnector( |
| 392 | media_io_kwargs={ |
| 393 | "video": { |
| 394 | "num_frames": num_frames, |
| 395 | } |
| 396 | }, |
| 397 | allowed_media_domains=[ |
| 398 | VLLM_S3_BUCKET_URL.removeprefix("https://"), |
| 399 | ], |
| 400 | ) |
| 401 | |
| 402 | try: |
| 403 | video_sync, metadata_sync = connector.fetch_video(video_url) |
| 404 | video_async, metadata_async = await connector.fetch_video_async(video_url) |
| 405 | except (TimeoutError, asyncio.TimeoutError) as e: |
| 406 | pytest.skip(f"Timeout fetching video (CI network flakiness): {e}") |
| 407 | |
| 408 | assert np.array_equal(video_sync, video_async) |
| 409 | assert metadata_sync == metadata_async |
| 410 | |
| 411 | disallowed_url = "https://upload.wikimedia.org/wikipedia/commons/4/47/PNG_transparency_demonstration_1.png" |
| 412 | with pytest.raises(ValueError): |
| 413 | _, _ = connector.fetch_video(disallowed_url) |
| 414 | |
| 415 | with pytest.raises(ValueError): |
| 416 | _, _ = await connector.fetch_video_async(disallowed_url) |
| 417 | |
| 418 | |
| 419 | @pytest.mark.asyncio |
| 420 | async def test_ssrf_bypass_backslash_in_url(local_asset_server): |
| 421 | """Verify that backslash-@ URL parsing confusion cannot bypass the |
| 422 | allowed_media_domains check (GHSA-v359-jj2v-j536). |
| 423 | |
| 424 | urllib3.parse_url() and aiohttp/yarl disagree on how to parse a |
| 425 | backslash before ``@``. urllib3 treats ``\\`` as part of the path |
| 426 | (encoding it as ``%5C``), while yarl treats it as a userinfo |
| 427 | separator, changing the effective host. The fix normalises the URL |
| 428 | through urllib3 *before* handing it to aiohttp so both layers agree. |
| 429 | """ |
| 430 | port = local_asset_server.port |
| 431 | asset = TEST_IMAGE_ASSETS[0] |
| 432 | |
| 433 | # Craft the bypass payload: urllib3 sees host=127.0.0.1, but an |
| 434 | # un-patched aiohttp would see host=example.com. |
| 435 | bypass_url = f"http://127.0.0.1:{port}\\@example.com/{asset}" |
| 436 | |
| 437 | connector = MediaConnector( |
| 438 | allowed_media_domains=["127.0.0.1"], |
| 439 | ) |
| 440 | |
| 441 | # After the fix the request is made to 127.0.0.1 (the local asset |
| 442 | # server) using the normalised URL. The normalised path will be |
| 443 | # /%5C@example.com/<asset> which won't match any file the server |
| 444 | # knows about, so we expect an HTTP error — but crucially NOT a |
| 445 | # successful fetch from example.com. |
| 446 | with pytest.raises(requests.exceptions.HTTPError): |
| 447 | connector.fetch_image(bypass_url) |
| 448 | |
| 449 | with pytest.raises(aiohttp.ClientResponseError): |
| 450 | await connector.fetch_image_async(bypass_url) |
| 451 | |
| 452 | |
| 453 | @pytest.mark.asyncio |
| 454 | async def test_ssrf_bypass_backslash_disallowed_domain(): |
| 455 | """The reverse direction: even when the *attacker-controlled* host |
| 456 | appears in the urllib3-parsed hostname position the allowlist must |
| 457 | still block it. |
| 458 | """ |
| 459 | # urllib3.parse_url sees host=example.com which is NOT in the |
| 460 | # allowlist, so this must be rejected before any request is made. |
| 461 | bypass_url = "https://example.com\\@safe.example.org/image.png" |
| 462 | |
| 463 | connector = MediaConnector( |
| 464 | allowed_media_domains=["safe.example.org"], |
| 465 | ) |
| 466 | |
| 467 | with pytest.raises(ValueError, match="allowed domains"): |
| 468 | connector.fetch_image(bypass_url) |
| 469 | |
| 470 | with pytest.raises(ValueError, match="allowed domains"): |
| 471 | await connector.fetch_image_async(bypass_url) |
| 472 | |
| 473 | |
| 474 | def _make_cached_connector(cache_dir, *, max_mb=10, ttl_hours=24): |
| 475 | """Create a MediaConnector with caching enabled via monkeypatched internals. |
| 476 | |
| 477 | We bypass __init__'s env-var path and wire up the cache fields directly |
| 478 | so tests don't depend on environment variables. URLs in these tests are |
| 479 | only used as cache keys (hashed to derive filenames); no HTTP requests |
| 480 | are made. |
| 481 | """ |
| 482 | connector = MediaConnector() |
| 483 | connector._media_cache_dir = cache_dir |
| 484 | connector._media_cache_max_bytes = max_mb * 1024 * 1024 |
| 485 | connector._media_cache_ttl_secs = ttl_hours * 3600 |
| 486 | return connector |
| 487 | |
| 488 | |
| 489 | def test_cache_put_and_get(): |
| 490 | """Basic round-trip: put bytes, get them back.""" |
| 491 | with TemporaryDirectory() as cache_dir: |
| 492 | connector = _make_cached_connector(cache_dir) |
| 493 | url = "https://example.com/image.png" |
| 494 | data = b"fake-image-bytes" |
| 495 | |
| 496 | connector._put_cached_bytes(url, data) |
| 497 | cached = connector._get_cached_bytes(url) |
| 498 | assert cached == data |
| 499 | |
| 500 | |
| 501 | def test_cache_ttl_expiry(): |
| 502 | """Entries older than TTL are evicted on read.""" |
| 503 | with TemporaryDirectory() as cache_dir: |
| 504 | connector = _make_cached_connector(cache_dir, ttl_hours=24) |
| 505 | url = "https://example.com/old.png" |
| 506 | data = b"old-data" |
| 507 | |
| 508 | connector._put_cached_bytes(url, data) |
| 509 | |
| 510 | # Backdate the file's mtime so it appears expired |
| 511 | cache_path = connector._media_cache_path(url) |
| 512 | expired_time = time.time() - (25 * 3600) # 25 hours ago |
| 513 | os.utime(cache_path, (expired_time, expired_time)) |
| 514 | |
| 515 | assert connector._get_cached_bytes(url) is None |
| 516 | assert not cache_path.exists() |
| 517 | |
| 518 | |
| 519 | def test_cache_lru_eviction(): |
| 520 | """Oldest entries are evicted when cache exceeds size budget.""" |
| 521 | with TemporaryDirectory() as cache_dir: |
| 522 | # Set a very small max size: 100 bytes |
| 523 | connector = _make_cached_connector(cache_dir, max_mb=0) |
| 524 | connector._media_cache_max_bytes = 100 |
| 525 | |
| 526 | # Write three 50-byte entries (total 150 > 100 budget) |
| 527 | urls = [f"https://example.com/{i}.png" for i in range(3)] |
| 528 | for i, url in enumerate(urls): |
| 529 | connector._put_cached_bytes(url, b"x" * 50) |
| 530 | # Stagger mtime so eviction order is deterministic |
| 531 | path = connector._media_cache_path(url) |
| 532 | os.utime(path, (time.time() + i, time.time() + i)) |
| 533 | |
| 534 | # The oldest entry (urls[0]) should have been evicted |
| 535 | assert connector._get_cached_bytes(urls[0]) is None |
| 536 | # The newest entries should still be present |
| 537 | assert connector._get_cached_bytes(urls[2]) == b"x" * 50 |
| 538 | |
| 539 | |
| 540 | def test_cache_ttl_eviction_during_write(): |
| 541 | """_maybe_evict removes expired files even if under size budget.""" |
| 542 | with TemporaryDirectory() as cache_dir: |
| 543 | connector = _make_cached_connector(cache_dir, ttl_hours=1) |
| 544 | url_old = "https://example.com/stale.png" |
| 545 | url_new = "https://example.com/fresh.png" |
| 546 | |
| 547 | connector._put_cached_bytes(url_old, b"stale") |
| 548 | # Backdate old entry past TTL |
| 549 | old_path = connector._media_cache_path(url_old) |
| 550 | expired_time = time.time() - (2 * 3600) |
| 551 | os.utime(old_path, (expired_time, expired_time)) |
| 552 | |
| 553 | # Writing a new entry triggers _maybe_evict |
| 554 | connector._put_cached_bytes(url_new, b"fresh") |
| 555 | |
| 556 | assert not old_path.exists() |
| 557 | assert connector._get_cached_bytes(url_new) == b"fresh" |
| 558 | |
| 559 | |
| 560 | def test_put_cached_bytes_missing_dir(): |
| 561 | """_put_cached_bytes does not crash when the cache dir disappears.""" |
| 562 | with TemporaryDirectory() as cache_dir: |
| 563 | connector = _make_cached_connector(cache_dir) |
| 564 | # Remove the directory to simulate it disappearing at runtime |
| 565 | shutil.rmtree(cache_dir) |
| 566 | |
| 567 | # Should not raise (graceful degradation) |
| 568 | connector._put_cached_bytes("https://example.com/x.png", b"data") |
| 569 | |
| 570 | |
| 571 | def test_get_cached_bytes_file_deleted_before_read(): |
| 572 | """_get_cached_bytes returns None if the file vanishes mid-read.""" |
| 573 | with TemporaryDirectory() as cache_dir: |
| 574 | connector = _make_cached_connector(cache_dir) |
| 575 | url = "https://example.com/vanish.png" |
| 576 | |
| 577 | connector._put_cached_bytes(url, b"data") |
| 578 | # Delete the file to simulate concurrent eviction |
| 579 | connector._media_cache_path(url).unlink() |
| 580 | |
| 581 | assert connector._get_cached_bytes(url) is None |