1
+ import os
1
2
from unittest .mock import MagicMock , patch
2
3
3
4
import pytest
4
5
5
6
from autogen .agentchat .conversable_agent import ConversableAgent
6
7
7
8
try :
9
+ from PIL import Image
10
+
8
11
from autogen .agentchat .contrib .capabilities .vision_capability import VisionCapability
9
12
except ImportError :
10
13
skip_test = True
@@ -21,6 +24,15 @@ def lmm_config():
21
24
}
22
25
23
26
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
+
24
36
@pytest .fixture
25
37
def vision_capability (lmm_config ):
26
38
return VisionCapability (lmm_config , custom_caption_func = None )
@@ -72,9 +84,9 @@ def test_process_last_received_message_text(mock_lmm_client, vision_capability):
72
84
def test_process_last_received_message_with_image (
73
85
mock_get_caption , mock_convert_base64 , mock_get_image_data , vision_capability
74
86
):
75
- content = [{"type" : "image_url" , "image_url" : {"url" : "notebook/viz_gc.png" }}]
87
+ content = [{"type" : "image_url" , "image_url" : {"url" : ( png_filename ()) }}]
76
88
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 "
78
90
)
79
91
processed_content = vision_capability .process_last_received_message (content )
80
92
assert processed_content == expected_caption
@@ -101,15 +113,15 @@ def caption_func(image_url: str, image_data=None, lmm_client=None) -> str:
101
113
class TestCustomCaptionFunc :
102
114
def test_custom_caption_func_with_valid_url (self , custom_caption_func ):
103
115
"""Test custom caption function with a valid image URL."""
104
- image_url = "notebook/viz_gc.png"
116
+ image_url = png_filename ()
105
117
expected_caption = f"An image description. The image is from { image_url } ."
106
118
assert custom_caption_func (image_url ) == expected_caption , "Caption does not match expected output."
107
119
108
120
def test_process_last_received_message_with_custom_func (self , lmm_config , custom_caption_func ):
109
121
"""Test processing a message containing an image URL with a custom caption function."""
110
122
vision_capability = VisionCapability (lmm_config , custom_caption_func = custom_caption_func )
111
123
112
- image_url = "notebook/viz_gc.png"
124
+ image_url = png_filename ()
113
125
content = [{"type" : "image_url" , "image_url" : {"url" : image_url }}]
114
126
expected_output = f" An image description. The image is from { image_url } ."
115
127
processed_content = vision_capability .process_last_received_message (content )
0 commit comments