65 lines
2.7 KiB
Python
65 lines
2.7 KiB
Python
import json
|
|
import subprocess
|
|
from unittest.mock import patch, MagicMock
|
|
from src.gemini_cli_adapter import GeminiCliAdapter
|
|
from src import mcp_client
|
|
|
|
def test_gemini_cli_context_bleed_prevention(monkeypatch) -> None:
|
|
"""Test that the GeminiCliAdapter correctly filters out echoed 'user' messages
|
|
from the streaming JSON if they were to occur (safety check)."""
|
|
adapter = GeminiCliAdapter()
|
|
|
|
mock_process = MagicMock()
|
|
# Simulate a stream that includes a message from 'user' (should be ignored)
|
|
# and a message from 'model'.
|
|
mock_process.stdout = [
|
|
b'{"kind": "message", "role": "user", "payload": "Echoed user prompt"}\n',
|
|
b'{"kind": "message", "role": "model", "payload": "Model response"}\n'
|
|
]
|
|
mock_process.stderr = []
|
|
mock_process.returncode = 0
|
|
|
|
with patch('subprocess.Popen', return_value=mock_process):
|
|
result = adapter.send("msg")
|
|
# Should only contain the model response
|
|
assert result["text"] == "Model response"
|
|
|
|
def test_gemini_cli_parameter_resilience() -> None:
|
|
"""Test that mcp_client correctly handles 'file_path' and 'dir_path' aliases
|
|
if the AI provides them instead of 'path'."""
|
|
from src import mcp_client
|
|
|
|
# Mock dispatch to see what it receives
|
|
with patch('src.mcp_client.read_file', return_value="content") as mock_read:
|
|
mcp_client.dispatch("read_file", {"file_path": "aliased.txt"})
|
|
mock_read.assert_called_once_with("aliased.txt")
|
|
|
|
with patch('src.mcp_client.list_directory', return_value="files") as mock_list:
|
|
mcp_client.dispatch("list_directory", {"dir_path": "aliased_dir"})
|
|
mock_list.assert_called_once_with("aliased_dir")
|
|
|
|
def test_gemini_cli_loop_termination() -> None:
|
|
"""Test that multi-round tool calling correctly terminates and preserves
|
|
the final text."""
|
|
from src import ai_client
|
|
|
|
ai_client.set_provider("gemini_cli", "gemini-2.0-flash")
|
|
|
|
# Round 1: Tool call
|
|
mock_resp1 = {"text": "Calling tool", "tool_calls": [{"name": "read_file", "args": {"path": "f.txt"}}]}
|
|
# Round 2: Final response
|
|
mock_resp2 = {"text": "Final answer", "tool_calls": []}
|
|
|
|
with patch('src.ai_client.GeminiCliAdapter') as MockAdapter:
|
|
instance = MockAdapter.return_value
|
|
instance.send.side_effect = [mock_resp1, mock_resp2]
|
|
instance.last_usage = {"total_tokens": 10}
|
|
instance.last_latency = 0.1
|
|
instance.session_id = "s1"
|
|
|
|
# We need to mock mcp_client.dispatch too
|
|
with patch('src.mcp_client.dispatch', return_value="content"):
|
|
result = ai_client.send("context", "prompt")
|
|
assert result == "Final answer"
|
|
assert instance.send.call_count == 2
|