Allegro
10271c6b44
Replace shell=True with list-based subprocess execution to prevent command injection via malicious user input. Changes: - tools/transcription_tools.py: Use shlex.split() + shell=False - tools/environments/docker.py: List-based commands with container ID validation Fixes CVE-level vulnerability where malicious file paths or container IDs could inject arbitrary commands. CVSS: 9.8 (Critical) Refs: V-001 in SECURITY_AUDIT_REPORT.md
308 lines
10 KiB
Python
308 lines
10 KiB
Python
"""Tests for agent/gemini_adapter.py - Google Gemini model support.
|
|
|
|
Tests message conversion, tool formatting, and response normalization.
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import patch, MagicMock
|
|
from types import SimpleNamespace
|
|
|
|
try:
|
|
from agent.gemini_adapter import (
|
|
convert_messages_to_gemini,
|
|
convert_tools_to_gemini,
|
|
normalize_gemini_response,
|
|
build_gemini_client,
|
|
GEMINI_ROLES,
|
|
)
|
|
HAS_MODULE = True
|
|
except ImportError:
|
|
HAS_MODULE = False
|
|
|
|
|
|
pytestmark = pytest.mark.skipif(not HAS_MODULE, reason="gemini_adapter module not found")
|
|
|
|
|
|
class TestConvertMessagesToGemini:
|
|
"""Tests for message format conversion."""
|
|
|
|
def test_converts_simple_user_message(self):
|
|
"""Should convert simple user message to Gemini format."""
|
|
messages = [{"role": "user", "content": "Hello"}]
|
|
result = convert_messages_to_gemini(messages)
|
|
|
|
assert len(result) == 1
|
|
assert result[0]["role"] == "user"
|
|
assert result[0]["parts"][0]["text"] == "Hello"
|
|
|
|
def test_converts_assistant_message(self):
|
|
"""Should convert assistant message to Gemini format."""
|
|
messages = [{"role": "assistant", "content": "Hi there!"}]
|
|
result = convert_messages_to_gemini(messages)
|
|
|
|
assert result[0]["role"] == "model"
|
|
assert result[0]["parts"][0]["text"] == "Hi there!"
|
|
|
|
def test_converts_system_message(self):
|
|
"""Should convert system message to Gemini format."""
|
|
messages = [{"role": "system", "content": "You are a helpful assistant."}]
|
|
result = convert_messages_to_gemini(messages)
|
|
|
|
# Gemini uses "user" role for system in some versions
|
|
assert result[0]["role"] in ["user", "system"]
|
|
|
|
def test_converts_tool_call_message(self):
|
|
"""Should convert tool call message."""
|
|
messages = [{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [{
|
|
"id": "call_123",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"arguments": '{"location": "NYC"}'
|
|
}
|
|
}]
|
|
}]
|
|
result = convert_messages_to_gemini(messages)
|
|
|
|
assert "function_call" in str(result)
|
|
|
|
def test_converts_tool_result_message(self):
|
|
"""Should convert tool result message."""
|
|
messages = [{
|
|
"role": "tool",
|
|
"tool_call_id": "call_123",
|
|
"content": '{"temperature": 72}'
|
|
}]
|
|
result = convert_messages_to_gemini(messages)
|
|
|
|
assert len(result) == 1
|
|
|
|
def test_handles_multipart_content(self):
|
|
"""Should handle messages with text and images."""
|
|
messages = [{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "What's in this image?"},
|
|
{"type": "image_url", "image_url": {"url": "data:image/png;base64,abc123"}}
|
|
]
|
|
}]
|
|
result = convert_messages_to_gemini(messages)
|
|
|
|
# Should have both text and image parts
|
|
parts = result[0]["parts"]
|
|
assert any(p.get("text") for p in parts)
|
|
assert any(p.get("inline_data") for p in parts)
|
|
|
|
|
|
class TestConvertToolsToGemini:
|
|
"""Tests for tool schema conversion."""
|
|
|
|
def test_converts_simple_function(self):
|
|
"""Should convert simple function tool."""
|
|
tools = [{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"description": "Get weather for a location",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"location": {"type": "string"}
|
|
},
|
|
"required": ["location"]
|
|
}
|
|
}
|
|
}]
|
|
result = convert_tools_to_gemini(tools)
|
|
|
|
assert len(result) == 1
|
|
assert result[0]["name"] == "get_weather"
|
|
assert "description" in result[0]
|
|
|
|
def test_converts_multiple_tools(self):
|
|
"""Should convert multiple tools."""
|
|
tools = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "tool_a",
|
|
"description": "Tool A",
|
|
"parameters": {"type": "object", "properties": {}}
|
|
}
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "tool_b",
|
|
"description": "Tool B",
|
|
"parameters": {"type": "object", "properties": {}}
|
|
}
|
|
}
|
|
]
|
|
result = convert_tools_to_gemini(tools)
|
|
|
|
assert len(result) == 2
|
|
assert result[0]["name"] == "tool_a"
|
|
assert result[1]["name"] == "tool_b"
|
|
|
|
def test_handles_complex_parameters(self):
|
|
"""Should handle complex parameter schemas."""
|
|
tools = [{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "complex_tool",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"count": {"type": "integer", "minimum": 0},
|
|
"items": {
|
|
"type": "array",
|
|
"items": {"type": "string"}
|
|
},
|
|
"config": {
|
|
"type": "object",
|
|
"properties": {
|
|
"enabled": {"type": "boolean"}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}]
|
|
result = convert_tools_to_gemini(tools)
|
|
|
|
assert result[0]["name"] == "complex_tool"
|
|
|
|
|
|
class TestNormalizeGeminiResponse:
|
|
"""Tests for response normalization."""
|
|
|
|
def test_normalizes_simple_text_response(self):
|
|
"""Should normalize simple text response."""
|
|
gemini_response = SimpleNamespace(
|
|
candidates=[SimpleNamespace(
|
|
content=SimpleNamespace(
|
|
parts=[SimpleNamespace(text="Hello!")]
|
|
),
|
|
finish_reason="STOP"
|
|
)]
|
|
)
|
|
result = normalize_gemini_response(gemini_response)
|
|
|
|
assert result.choices[0].message.content == "Hello!"
|
|
assert result.choices[0].finish_reason == "stop"
|
|
|
|
def test_normalizes_tool_call_response(self):
|
|
"""Should normalize tool call response."""
|
|
gemini_response = SimpleNamespace(
|
|
candidates=[SimpleNamespace(
|
|
content=SimpleNamespace(
|
|
parts=[SimpleNamespace(
|
|
function_call=SimpleNamespace(
|
|
name="get_weather",
|
|
args={"location": "NYC"}
|
|
)
|
|
)]
|
|
),
|
|
finish_reason="STOP"
|
|
)]
|
|
)
|
|
result = normalize_gemini_response(gemini_response)
|
|
|
|
assert result.choices[0].message.tool_calls is not None
|
|
assert result.choices[0].message.tool_calls[0].function.name == "get_weather"
|
|
|
|
def test_handles_empty_response(self):
|
|
"""Should handle empty response gracefully."""
|
|
gemini_response = SimpleNamespace(
|
|
candidates=[SimpleNamespace(
|
|
content=SimpleNamespace(parts=[]),
|
|
finish_reason="STOP"
|
|
)]
|
|
)
|
|
result = normalize_gemini_response(gemini_response)
|
|
|
|
assert result.choices[0].message.content == ""
|
|
|
|
def test_handles_safety_blocked_response(self):
|
|
"""Should handle safety-blocked response."""
|
|
gemini_response = SimpleNamespace(
|
|
candidates=[SimpleNamespace(
|
|
finish_reason="SAFETY",
|
|
safety_ratings=[SimpleNamespace(
|
|
category="HARM_CATEGORY_DANGEROUS_CONTENT",
|
|
probability="HIGH"
|
|
)]
|
|
)]
|
|
)
|
|
result = normalize_gemini_response(gemini_response)
|
|
|
|
assert result.choices[0].finish_reason == "content_filter"
|
|
|
|
def test_extracts_usage_info(self):
|
|
"""Should extract token usage if available."""
|
|
gemini_response = SimpleNamespace(
|
|
candidates=[SimpleNamespace(
|
|
content=SimpleNamespace(parts=[SimpleNamespace(text="Hi")]),
|
|
finish_reason="STOP"
|
|
)],
|
|
usage_metadata=SimpleNamespace(
|
|
prompt_token_count=10,
|
|
candidates_token_count=5,
|
|
total_token_count=15
|
|
)
|
|
)
|
|
result = normalize_gemini_response(gemini_response)
|
|
|
|
assert result.usage.prompt_tokens == 10
|
|
assert result.usage.completion_tokens == 5
|
|
assert result.usage.total_tokens == 15
|
|
|
|
|
|
class TestBuildGeminiClient:
|
|
"""Tests for client initialization."""
|
|
|
|
def test_builds_client_with_api_key(self):
|
|
"""Should build client with API key."""
|
|
with patch("agent.gemini_adapter.genai") as mock_genai:
|
|
mock_client = MagicMock()
|
|
mock_genai.GenerativeModel.return_value = mock_client
|
|
|
|
client = build_gemini_client(api_key="test-key-123")
|
|
|
|
mock_genai.configure.assert_called_once_with(api_key="test-key-123")
|
|
|
|
def test_applies_generation_config(self):
|
|
"""Should apply generation configuration."""
|
|
with patch("agent.gemini_adapter.genai") as mock_genai:
|
|
build_gemini_client(
|
|
api_key="test-key",
|
|
temperature=0.5,
|
|
max_output_tokens=1000,
|
|
top_p=0.9
|
|
)
|
|
|
|
call_kwargs = mock_genai.GenerativeModel.call_args[1]
|
|
assert "generation_config" in call_kwargs
|
|
|
|
|
|
class TestGeminiRoleMapping:
|
|
"""Tests for role mapping between OpenAI and Gemini formats."""
|
|
|
|
def test_user_role_mapping(self):
|
|
"""Should map user role correctly."""
|
|
assert "user" in GEMINI_ROLES.values() or "user" in str(GEMINI_ROLES)
|
|
|
|
def test_assistant_role_mapping(self):
|
|
"""Should map assistant to model role."""
|
|
# Gemini uses "model" instead of "assistant"
|
|
assert GEMINI_ROLES.get("assistant") == "model" or "model" in str(GEMINI_ROLES)
|
|
|
|
def test_system_role_mapping(self):
|
|
"""Should handle system role appropriately."""
|
|
# System messages handled differently in Gemini
|
|
assert "system" in str(GEMINI_ROLES).lower() or True # Implementation dependent
|