security: fix command injection vulnerabilities (CVSS 9.8)
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
This commit is contained in:
307
tests/agent/test_gemini_adapter.py
Normal file
307
tests/agent/test_gemini_adapter.py
Normal file
@@ -0,0 +1,307 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user