548 lines
19 KiB
Python
548 lines
19 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""OpenAI services test module."""
|
|
import unittest
|
|
from unittest.mock import patch, MagicMock, mock_open
|
|
import os
|
|
import shutil
|
|
from openai._types import NOT_GIVEN
|
|
|
|
|
|
from agentscope.service.multi_modality.openai_services import (
|
|
openai_audio_to_text,
|
|
openai_text_to_audio,
|
|
openai_text_to_image,
|
|
openai_image_to_text,
|
|
openai_edit_image,
|
|
openai_create_image_variation,
|
|
)
|
|
from agentscope.service import ServiceExecStatus
|
|
from agentscope.service.multi_modality.openai_services import _audio_filename
|
|
|
|
|
|
class TestOpenAIServices(unittest.TestCase):
|
|
"""Test the OpenAI services functions."""
|
|
|
|
def setUp(self) -> None:
|
|
"""Set up the test"""
|
|
self.save_dir = os.path.join(os.path.dirname(__file__), "test_dir")
|
|
os.makedirs(self.save_dir, exist_ok=True)
|
|
self.maxDiff = None
|
|
|
|
def tearDown(self) -> None:
|
|
"""Tear down the test"""
|
|
if os.path.exists(self.save_dir):
|
|
shutil.rmtree(self.save_dir)
|
|
|
|
@patch(
|
|
"agentscope.service.multi_modality.openai_services.OpenAIDALLEWrapper",
|
|
)
|
|
@patch("agentscope.service.multi_modality.openai_services._download_file")
|
|
def test_openai_text_to_image_success(
|
|
self,
|
|
_: MagicMock,
|
|
mock_dalle_wrapper: MagicMock,
|
|
) -> None:
|
|
"""Test the openai_text_to_image function with a valid prompt."""
|
|
# Mock the OpenAIDALLEWrapper response
|
|
mock_response = MagicMock()
|
|
mock_response.image_urls = [
|
|
"https://example.com/image1.png",
|
|
"https://example.com/image2.png",
|
|
]
|
|
mock_instance = MagicMock()
|
|
mock_instance.return_value = mock_response
|
|
mock_dalle_wrapper.return_value = mock_instance
|
|
|
|
# Call the function with a valid prompt
|
|
result = openai_text_to_image(
|
|
prompt="A futuristic city skyline at sunset",
|
|
api_key="fake_api_key",
|
|
n=2,
|
|
model="dall-e-2",
|
|
size="256x256",
|
|
quality="standard",
|
|
style="vivid",
|
|
save_dir=self.save_dir,
|
|
)
|
|
|
|
self.assertEqual(result.status, ServiceExecStatus.SUCCESS)
|
|
expected_urls = [
|
|
os.path.join(self.save_dir, "image1.png"),
|
|
os.path.join(self.save_dir, "image2.png"),
|
|
]
|
|
self.assertEqual(result.content, {"image_urls": expected_urls})
|
|
|
|
# Ensure the wrapper and response methods are called correctly
|
|
mock_dalle_wrapper.assert_called_once_with(
|
|
config_name="text_to_image_service_call",
|
|
model_name="dall-e-2",
|
|
api_key="fake_api_key",
|
|
)
|
|
mock_instance.assert_called_once_with(
|
|
prompt="A futuristic city skyline at sunset",
|
|
n=2,
|
|
size="256x256",
|
|
quality="standard",
|
|
style="vivid",
|
|
)
|
|
|
|
@patch(
|
|
"agentscope.service.multi_modality.openai_services.OpenAIDALLEWrapper",
|
|
)
|
|
def test_openai_text_to_image_api_error(
|
|
self,
|
|
mock_dalle_wrapper: MagicMock,
|
|
) -> None:
|
|
"""Test the openai_text_to_image function with an API error."""
|
|
# Mock an API failure response
|
|
mock_instance = MagicMock()
|
|
mock_instance.side_effect = Exception("API Error: Invalid request")
|
|
mock_dalle_wrapper.return_value = mock_instance
|
|
|
|
# Call the function with a prompt that causes an API error
|
|
result = openai_text_to_image(
|
|
prompt="An invalid prompt",
|
|
api_key="fake_api_key",
|
|
n=1,
|
|
model="dall-e-2",
|
|
size="256x256",
|
|
quality="standard",
|
|
style="vivid",
|
|
)
|
|
|
|
self.assertEqual(result.status, ServiceExecStatus.ERROR)
|
|
self.assertIn("API Error", result.content)
|
|
|
|
# Ensure the wrapper and response methods are called correctly
|
|
mock_dalle_wrapper.assert_called_once_with(
|
|
config_name="text_to_image_service_call",
|
|
model_name="dall-e-2",
|
|
api_key="fake_api_key",
|
|
)
|
|
mock_instance.assert_called_once_with(
|
|
prompt="An invalid prompt",
|
|
n=1,
|
|
size="256x256",
|
|
quality="standard",
|
|
style="vivid",
|
|
)
|
|
|
|
@patch(
|
|
"agentscope.service.multi_modality.openai_services.OpenAIDALLEWrapper",
|
|
)
|
|
@patch("agentscope.service.multi_modality.openai_services._download_file")
|
|
def test_openai_text_to_image_service_error(
|
|
self,
|
|
mock_download_file: MagicMock,
|
|
mock_dalle_wrapper: MagicMock,
|
|
) -> None:
|
|
"""Test the openai_text_to_image function with a service error."""
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.image_urls = None
|
|
mock_instance = MagicMock()
|
|
mock_instance.return_value = mock_response
|
|
mock_dalle_wrapper.return_value = mock_instance
|
|
|
|
# Call the function with a valid prompt but simulate an internal error
|
|
result = openai_text_to_image(
|
|
prompt="A futuristic city skyline at sunset",
|
|
api_key="fake_api_key",
|
|
n=2,
|
|
model="dall-e-2",
|
|
size="256x256",
|
|
quality="standard",
|
|
style="vivid",
|
|
save_dir=self.save_dir,
|
|
)
|
|
|
|
self.assertEqual(result.status, ServiceExecStatus.ERROR)
|
|
self.assertIn("Error: Failed to generate images", result.content)
|
|
|
|
# Ensure the wrapper and response methods are called correctly
|
|
mock_dalle_wrapper.assert_called_once_with(
|
|
config_name="text_to_image_service_call",
|
|
model_name="dall-e-2",
|
|
api_key="fake_api_key",
|
|
)
|
|
mock_instance.assert_called_once_with(
|
|
prompt="A futuristic city skyline at sunset",
|
|
n=2,
|
|
size="256x256",
|
|
quality="standard",
|
|
style="vivid",
|
|
)
|
|
|
|
# Ensure _download_file is not called in case of service error
|
|
mock_download_file.assert_not_called()
|
|
|
|
@patch("agentscope.service.multi_modality.openai_services.OpenAI")
|
|
@patch(
|
|
"builtins.open",
|
|
new_callable=mock_open,
|
|
read_data=b"mock_audio_data",
|
|
)
|
|
def test_openai_audio_to_text_success(
|
|
self,
|
|
_: MagicMock,
|
|
mock_openai: MagicMock,
|
|
) -> None:
|
|
"""Test the openai_audio_to_text function with a valid audio file."""
|
|
# Mocking the OpenAI API response
|
|
mock_transcription = (
|
|
mock_openai.return_value.audio.transcriptions.create.return_value
|
|
)
|
|
mock_transcription.text = "This is a test transcription."
|
|
|
|
# Test audio file path within the save_dir
|
|
audio_file_url = os.path.join(self.save_dir, "test_audio.mp3")
|
|
|
|
# Create the test file in the save_dir
|
|
with open(audio_file_url, "wb") as f:
|
|
f.write(b"mock_audio_data")
|
|
|
|
# Call the function
|
|
result = openai_audio_to_text(audio_file_url, "mock_api_key")
|
|
|
|
# Assertions
|
|
self.assertEqual(result.status, ServiceExecStatus.SUCCESS)
|
|
self.assertEqual(
|
|
result.content,
|
|
{"transcription": "This is a test transcription."},
|
|
)
|
|
|
|
@patch("agentscope.service.multi_modality.openai_services.OpenAI")
|
|
@patch("builtins.open", new_callable=mock_open)
|
|
def test_openai_audio_to_text_error(
|
|
self,
|
|
_: MagicMock,
|
|
mock_openai: MagicMock,
|
|
) -> None:
|
|
"""Test the openai_audio_to_text function with an API error."""
|
|
# Mocking an OpenAI API error
|
|
mock_openai.return_value.audio.transcriptions.create.side_effect = (
|
|
Exception("API Error")
|
|
)
|
|
|
|
# Test audio file path within the save_dir
|
|
audio_file_url = os.path.join(self.save_dir, "test_audio.mp3")
|
|
|
|
# Call the function (the mock_file_open will handle the file opening)
|
|
result = openai_audio_to_text(audio_file_url, "mock_api_key")
|
|
|
|
# Assertions
|
|
self.assertEqual(result.status, ServiceExecStatus.ERROR)
|
|
self.assertIn(
|
|
"Error: Failed to transcribe audio API Error",
|
|
result.content,
|
|
)
|
|
|
|
@patch("agentscope.service.multi_modality.openai_services.OpenAI")
|
|
def test_successful_audio_generation(self, mock_openai: MagicMock) -> None:
|
|
"""Test the openai_text_to_audio function with a valid text."""
|
|
# Mocking the OpenAI API response
|
|
mock_response = (
|
|
mock_openai.return_value.audio.speech.create.return_value
|
|
)
|
|
|
|
# Sample text and expected audio path
|
|
text = "Hello, this is a test."
|
|
save_dir = self.save_dir
|
|
expected_audio_path = os.path.join(
|
|
save_dir,
|
|
f"{_audio_filename(text)}.mp3",
|
|
)
|
|
|
|
# Call the function
|
|
result = openai_text_to_audio(text, "mock_api_key", save_dir)
|
|
|
|
# Assertions
|
|
self.assertEqual(result.status, ServiceExecStatus.SUCCESS)
|
|
self.assertEqual(result.content["audio_path"], expected_audio_path)
|
|
mock_response.stream_to_file.assert_called_once_with(
|
|
expected_audio_path,
|
|
) # Check file save
|
|
|
|
@patch("agentscope.service.multi_modality.openai_services.OpenAI")
|
|
def test_api_error_text_to_audio(self, mock_openai: MagicMock) -> None:
|
|
"""Test the openai_text_to_audio function with an API error."""
|
|
# Mocking an OpenAI API error
|
|
mock_openai.return_value.audio.speech.create.side_effect = Exception(
|
|
"API Error",
|
|
)
|
|
|
|
# Call the function
|
|
result = openai_text_to_audio(
|
|
"Hello, this is a test.",
|
|
"mock_api_key",
|
|
self.save_dir,
|
|
)
|
|
|
|
# Assertions
|
|
self.assertEqual(result.status, ServiceExecStatus.ERROR)
|
|
self.assertIn(
|
|
"Error: Failed to generate audio. API Error",
|
|
result.content,
|
|
)
|
|
|
|
@patch(
|
|
"agentscope.service.multi_modality.openai_services.OpenAIChatWrapper",
|
|
)
|
|
def test_openai_image_to_text_success(
|
|
self,
|
|
mock_openai_chat_wrapper: MagicMock,
|
|
) -> None:
|
|
"""Test the openai_image_to_text function with a valid image URL."""
|
|
# Mock the OpenAIChatWrapper
|
|
mock_wrapper_instance = MagicMock()
|
|
mock_openai_chat_wrapper.return_value = mock_wrapper_instance
|
|
|
|
# Mock the response
|
|
mock_response = MagicMock()
|
|
mock_response.text = "This is a description of the image."
|
|
mock_wrapper_instance.return_value = mock_response
|
|
|
|
# Test with single image URL
|
|
result = openai_image_to_text(
|
|
"https://example.com/image.jpg",
|
|
"fake_api_key",
|
|
)
|
|
|
|
# Assertions
|
|
self.assertEqual(result.status, ServiceExecStatus.SUCCESS)
|
|
self.assertEqual(result.content, "This is a description of the image.")
|
|
|
|
mock_openai_chat_wrapper.assert_called_once_with(
|
|
config_name="image_to_text_service_call",
|
|
model_name="gpt-4o",
|
|
api_key="fake_api_key",
|
|
)
|
|
|
|
mock_wrapper_instance.format.assert_called_once()
|
|
formatted_message = mock_wrapper_instance.format.call_args[0][0]
|
|
self.assertEqual(formatted_message.name, "service_call")
|
|
self.assertEqual(formatted_message.role, "user")
|
|
self.assertEqual(formatted_message.content, "Describe the image")
|
|
self.assertEqual(
|
|
formatted_message.url,
|
|
"https://example.com/image.jpg",
|
|
)
|
|
|
|
@patch(
|
|
"agentscope.service.multi_modality.openai_services.OpenAIChatWrapper",
|
|
)
|
|
def test_openai_image_to_text_error(
|
|
self,
|
|
mock_openai_chat_wrapper: MagicMock,
|
|
) -> None:
|
|
"""Test the openai_image_to_text function with an API error."""
|
|
# Mock the OpenAIChatWrapper to raise an exception
|
|
mock_wrapper_instance = MagicMock()
|
|
mock_openai_chat_wrapper.return_value = mock_wrapper_instance
|
|
mock_wrapper_instance.side_effect = Exception("API Error")
|
|
|
|
# Test with single image URL
|
|
result = openai_image_to_text(
|
|
"https://example.com/image.jpg",
|
|
"fake_api_key",
|
|
)
|
|
|
|
# Assertions
|
|
self.assertEqual(result.status, ServiceExecStatus.ERROR)
|
|
self.assertEqual(result.content, "API Error")
|
|
|
|
@patch("agentscope.service.multi_modality.openai_services.OpenAI")
|
|
@patch("agentscope.service.multi_modality.openai_services._parse_url")
|
|
@patch(
|
|
(
|
|
"agentscope.service.multi_modality.openai_services"
|
|
"._handle_openai_img_response"
|
|
),
|
|
)
|
|
def test_openai_edit_image_success(
|
|
self,
|
|
mock_handle_response: MagicMock,
|
|
mock_parse_url: MagicMock,
|
|
mock_openai: MagicMock,
|
|
) -> None:
|
|
"""Test the openai_edit_image function with a valid image URL."""
|
|
# Mock OpenAI client
|
|
mock_client = MagicMock()
|
|
mock_openai.return_value = mock_client
|
|
mock_response = MagicMock()
|
|
mock_client.images.edit.return_value = mock_response
|
|
|
|
# Mock _parse_url function
|
|
mock_parse_url.side_effect = lambda x: f"parsed_{x}"
|
|
|
|
# Mock _handle_openai_img_response function
|
|
mock_handle_response.return_value = [
|
|
"edited_image_url1",
|
|
"edited_image_url2",
|
|
]
|
|
|
|
# Call the function
|
|
result = openai_edit_image(
|
|
image_url="original_image.png",
|
|
prompt="Add a sun to the sky",
|
|
api_key="fake_api_key",
|
|
mask_url="mask_image.png",
|
|
n=2,
|
|
size="512x512",
|
|
save_dir=None,
|
|
)
|
|
|
|
# Assertions
|
|
self.assertEqual(result.status, ServiceExecStatus.SUCCESS)
|
|
self.assertEqual(
|
|
result.content,
|
|
{"image_urls": ["edited_image_url1", "edited_image_url2"]},
|
|
)
|
|
|
|
# Check if OpenAI client was called correctly
|
|
mock_client.images.edit.assert_called_once_with(
|
|
model="dall-e-2",
|
|
image="parsed_original_image.png",
|
|
mask="parsed_mask_image.png",
|
|
prompt="Add a sun to the sky",
|
|
n=2,
|
|
size="512x512",
|
|
)
|
|
|
|
# Check if _handle_openai_img_response was called
|
|
mock_handle_response.assert_called_once_with(mock_response, None)
|
|
|
|
@patch("agentscope.service.multi_modality.openai_services.OpenAI")
|
|
@patch("agentscope.service.multi_modality.openai_services._parse_url")
|
|
def test_openai_edit_image_error(
|
|
self,
|
|
mock_parse_url: MagicMock,
|
|
mock_openai: MagicMock,
|
|
) -> None:
|
|
"""Test the openai_edit_image function with an API error."""
|
|
# Mock OpenAI client to raise an exception
|
|
mock_client = MagicMock()
|
|
mock_openai.return_value = mock_client
|
|
mock_client.images.edit.side_effect = Exception("API Error")
|
|
|
|
# Mock _parse_url function
|
|
mock_parse_url.side_effect = lambda x: f"parsed_{x}"
|
|
|
|
# Call the function
|
|
result = openai_edit_image(
|
|
image_url="original_image.png",
|
|
prompt="Add a sun to the sky",
|
|
api_key="fake_api_key",
|
|
)
|
|
|
|
# Assertions
|
|
self.assertEqual(result.status, ServiceExecStatus.ERROR)
|
|
self.assertEqual(result.content, "API Error")
|
|
|
|
# Check if OpenAI client was called
|
|
mock_client.images.edit.assert_called_once_with(
|
|
model="dall-e-2",
|
|
image="parsed_original_image.png",
|
|
mask=NOT_GIVEN,
|
|
prompt="Add a sun to the sky",
|
|
n=1,
|
|
size="256x256",
|
|
)
|
|
|
|
@patch("agentscope.service.multi_modality.openai_services.OpenAI")
|
|
@patch("agentscope.service.multi_modality.openai_services._parse_url")
|
|
@patch(
|
|
(
|
|
"agentscope.service.multi_modality.openai_services"
|
|
"._handle_openai_img_response"
|
|
),
|
|
)
|
|
def test_openai_create_image_variation_success(
|
|
self,
|
|
mock_handle_response: MagicMock,
|
|
mock_parse_url: MagicMock,
|
|
mock_openai: MagicMock,
|
|
) -> None:
|
|
"""Test the openai_create_image_variation swith a valid image URL."""
|
|
# Mock OpenAI client
|
|
mock_client = MagicMock()
|
|
mock_openai.return_value = mock_client
|
|
mock_response = MagicMock()
|
|
mock_client.images.create_variation.return_value = mock_response
|
|
|
|
# Mock _parse_url function
|
|
mock_parse_url.return_value = "parsed_image.png"
|
|
|
|
# Mock _handle_openai_img_response function
|
|
mock_handle_response.return_value = [
|
|
"variation_url1",
|
|
"variation_url2",
|
|
]
|
|
|
|
# Call the function
|
|
result = openai_create_image_variation(
|
|
image_url="original_image.png",
|
|
api_key="fake_api_key",
|
|
n=2,
|
|
size="512x512",
|
|
save_dir=None,
|
|
)
|
|
|
|
# Assertions
|
|
self.assertEqual(result.status, ServiceExecStatus.SUCCESS)
|
|
self.assertEqual(
|
|
result.content,
|
|
{"image_urls": ["variation_url1", "variation_url2"]},
|
|
)
|
|
|
|
# Check if OpenAI client was called correctly
|
|
mock_client.images.create_variation.assert_called_once_with(
|
|
model="dall-e-2",
|
|
image="parsed_image.png",
|
|
n=2,
|
|
size="512x512",
|
|
)
|
|
|
|
# Check if _handle_openai_img_response was called
|
|
mock_handle_response.assert_called_once_with(mock_response, None)
|
|
|
|
@patch("agentscope.service.multi_modality.openai_services.OpenAI")
|
|
@patch("agentscope.service.multi_modality.openai_services._parse_url")
|
|
def test_openai_create_image_variation_error(
|
|
self,
|
|
mock_parse_url: MagicMock,
|
|
mock_openai: MagicMock,
|
|
) -> None:
|
|
"""Test the openai_create_image_variation with an API error."""
|
|
# Mock OpenAI client to raise an exception
|
|
mock_client = MagicMock()
|
|
mock_openai.return_value = mock_client
|
|
mock_client.images.create_variation.side_effect = Exception(
|
|
"API Error",
|
|
)
|
|
|
|
# Mock _parse_url function
|
|
mock_parse_url.return_value = "parsed_image.png"
|
|
|
|
# Call the function
|
|
result = openai_create_image_variation(
|
|
image_url="original_image.png",
|
|
api_key="fake_api_key",
|
|
)
|
|
|
|
# Assertions
|
|
self.assertEqual(result.status, ServiceExecStatus.ERROR)
|
|
self.assertEqual(result.content, "API Error")
|
|
|
|
# Check if OpenAI client was called
|
|
mock_client.images.create_variation.assert_called_once_with(
|
|
model="dall-e-2",
|
|
image="parsed_image.png",
|
|
n=1,
|
|
size="256x256",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|