mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-28 14:07:55 +00:00
151 lines
5.3 KiB
Python
151 lines
5.3 KiB
Python
"""Tests for the image_generate tool."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import builtins
|
|
import sys
|
|
from unittest.mock import MagicMock
|
|
|
|
from openjarvis.tools.image_tool import ImageGenerateTool
|
|
|
|
|
|
class TestImageGenerateTool:
|
|
def test_spec(self):
|
|
tool = ImageGenerateTool()
|
|
assert tool.spec.name == "image_generate"
|
|
assert tool.spec.category == "media"
|
|
assert "prompt" in tool.spec.parameters["properties"]
|
|
assert "prompt" in tool.spec.parameters["required"]
|
|
assert tool.spec.required_capabilities == ["network:fetch"]
|
|
|
|
def test_tool_id(self):
|
|
tool = ImageGenerateTool()
|
|
assert tool.tool_id == "image_generate"
|
|
|
|
def test_no_prompt(self):
|
|
tool = ImageGenerateTool()
|
|
result = tool.execute(prompt="")
|
|
assert result.success is False
|
|
assert "No prompt" in result.content
|
|
|
|
def test_no_prompt_param(self):
|
|
tool = ImageGenerateTool()
|
|
result = tool.execute()
|
|
assert result.success is False
|
|
assert "No prompt" in result.content
|
|
|
|
def test_invalid_size(self):
|
|
tool = ImageGenerateTool()
|
|
result = tool.execute(prompt="a cat", size="999x999")
|
|
assert result.success is False
|
|
assert "Invalid size" in result.content
|
|
|
|
def test_unsupported_provider(self):
|
|
tool = ImageGenerateTool()
|
|
result = tool.execute(prompt="a cat", provider="midjourney")
|
|
assert result.success is False
|
|
assert "Unsupported provider" in result.content
|
|
|
|
def test_openai_not_installed(self, monkeypatch):
|
|
"""Simulate openai package not being installed."""
|
|
monkeypatch.delitem(sys.modules, "openai", raising=False)
|
|
original_import = builtins.__import__
|
|
|
|
def _mock_import(name, *args, **kwargs):
|
|
if name == "openai":
|
|
raise ImportError("No module named 'openai'")
|
|
return original_import(name, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(builtins, "__import__", _mock_import)
|
|
|
|
tool = ImageGenerateTool()
|
|
result = tool.execute(prompt="a cat")
|
|
assert result.success is False
|
|
assert "openai package not installed" in result.content
|
|
|
|
def test_no_api_key(self, monkeypatch):
|
|
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
|
mock_openai = MagicMock()
|
|
monkeypatch.setitem(sys.modules, "openai", mock_openai)
|
|
|
|
tool = ImageGenerateTool()
|
|
result = tool.execute(prompt="a cat")
|
|
assert result.success is False
|
|
assert "No API key" in result.content
|
|
|
|
def test_successful_generation(self, monkeypatch):
|
|
monkeypatch.setenv("OPENAI_API_KEY", "test-key")
|
|
mock_image_data = MagicMock()
|
|
mock_image_data.url = "https://example.com/image.png"
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.data = [mock_image_data]
|
|
|
|
mock_client = MagicMock()
|
|
mock_client.images.generate.return_value = mock_response
|
|
|
|
mock_openai = MagicMock()
|
|
mock_openai.OpenAI.return_value = mock_client
|
|
monkeypatch.setitem(sys.modules, "openai", mock_openai)
|
|
|
|
tool = ImageGenerateTool()
|
|
result = tool.execute(prompt="a cat on a mat")
|
|
assert result.success is True
|
|
assert result.content == "https://example.com/image.png"
|
|
assert result.metadata["url"] == "https://example.com/image.png"
|
|
assert result.metadata["size"] == "1024x1024"
|
|
assert result.metadata["provider"] == "openai"
|
|
|
|
def test_save_to_file(self, monkeypatch, tmp_path):
|
|
monkeypatch.setenv("OPENAI_API_KEY", "test-key")
|
|
mock_image_data = MagicMock()
|
|
mock_image_data.url = "https://example.com/image.png"
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.data = [mock_image_data]
|
|
|
|
mock_client = MagicMock()
|
|
mock_client.images.generate.return_value = mock_response
|
|
|
|
mock_openai = MagicMock()
|
|
mock_openai.OpenAI.return_value = mock_client
|
|
monkeypatch.setitem(sys.modules, "openai", mock_openai)
|
|
|
|
# Mock httpx for downloading
|
|
import httpx
|
|
|
|
mock_http_resp = MagicMock()
|
|
mock_http_resp.content = b"\x89PNG\r\n\x1a\nfake-image-data"
|
|
mock_http_resp.raise_for_status = MagicMock()
|
|
monkeypatch.setattr(httpx, "get", MagicMock(return_value=mock_http_resp))
|
|
|
|
output_file = tmp_path / "output.png"
|
|
tool = ImageGenerateTool()
|
|
result = tool.execute(
|
|
prompt="a cat",
|
|
output_path=str(output_file),
|
|
)
|
|
assert result.success is True
|
|
assert output_file.exists()
|
|
assert output_file.read_bytes() == b"\x89PNG\r\n\x1a\nfake-image-data"
|
|
|
|
def test_api_error(self, monkeypatch):
|
|
monkeypatch.setenv("OPENAI_API_KEY", "test-key")
|
|
mock_client = MagicMock()
|
|
mock_client.images.generate.side_effect = RuntimeError("Rate limit exceeded")
|
|
|
|
mock_openai = MagicMock()
|
|
mock_openai.OpenAI.return_value = mock_client
|
|
monkeypatch.setitem(sys.modules, "openai", mock_openai)
|
|
|
|
tool = ImageGenerateTool()
|
|
result = tool.execute(prompt="a cat")
|
|
assert result.success is False
|
|
assert "Image generation error" in result.content
|
|
|
|
def test_to_openai_function(self):
|
|
tool = ImageGenerateTool()
|
|
fn = tool.to_openai_function()
|
|
assert fn["type"] == "function"
|
|
assert fn["function"]["name"] == "image_generate"
|