Skip to content

Commit c030225

Browse files
authored
Update mm test: create dummy image in case file corrput (#2258)
1 parent c27eb0d commit c030225

File tree

1 file changed

+16
-4
lines changed

1 file changed

+16
-4
lines changed

test/agentchat/contrib/capabilities/test_vision_capability.py

+16-4
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,13 @@
1+
import os
12
from unittest.mock import MagicMock, patch
23

34
import pytest
45

56
from autogen.agentchat.conversable_agent import ConversableAgent
67

78
try:
9+
from PIL import Image
10+
811
from autogen.agentchat.contrib.capabilities.vision_capability import VisionCapability
912
except ImportError:
1013
skip_test = True
@@ -21,6 +24,15 @@ def lmm_config():
2124
}
2225

2326

27+
def png_filename() -> str:
28+
filename = "tmp/test_image.png"
29+
if not os.path.exists(filename):
30+
# Setup: Create a PNG file
31+
image = Image.new("RGB", (100, 100), color="blue")
32+
image.save(filename)
33+
return filename # This is what the test will use
34+
35+
2436
@pytest.fixture
2537
def vision_capability(lmm_config):
2638
return VisionCapability(lmm_config, custom_caption_func=None)
@@ -72,9 +84,9 @@ def test_process_last_received_message_text(mock_lmm_client, vision_capability):
7284
def test_process_last_received_message_with_image(
7385
mock_get_caption, mock_convert_base64, mock_get_image_data, vision_capability
7486
):
75-
content = [{"type": "image_url", "image_url": {"url": "notebook/viz_gc.png"}}]
87+
content = [{"type": "image_url", "image_url": {"url": (png_filename())}}]
7688
expected_caption = (
77-
"<img notebook/viz_gc.png> in case you can not see, the caption of this image is: A sample image caption.\n"
89+
f"<img {png_filename()}> in case you can not see, the caption of this image is: A sample image caption.\n"
7890
)
7991
processed_content = vision_capability.process_last_received_message(content)
8092
assert processed_content == expected_caption
@@ -101,15 +113,15 @@ def caption_func(image_url: str, image_data=None, lmm_client=None) -> str:
101113
class TestCustomCaptionFunc:
102114
def test_custom_caption_func_with_valid_url(self, custom_caption_func):
103115
"""Test custom caption function with a valid image URL."""
104-
image_url = "notebook/viz_gc.png"
116+
image_url = png_filename()
105117
expected_caption = f"An image description. The image is from {image_url}."
106118
assert custom_caption_func(image_url) == expected_caption, "Caption does not match expected output."
107119

108120
def test_process_last_received_message_with_custom_func(self, lmm_config, custom_caption_func):
109121
"""Test processing a message containing an image URL with a custom caption function."""
110122
vision_capability = VisionCapability(lmm_config, custom_caption_func=custom_caption_func)
111123

112-
image_url = "notebook/viz_gc.png"
124+
image_url = png_filename()
113125
content = [{"type": "image_url", "image_url": {"url": image_url}}]
114126
expected_output = f" An image description. The image is from {image_url}."
115127
processed_content = vision_capability.process_last_received_message(content)

0 commit comments

Comments
 (0)