ed
6a3c142bfa
Removed 18 unused 'from src import models' imports. Migrated the explicit sub-imports: - 'from src.models import PROVIDERS' -> 'from src.ai_client import PROVIDERS' (test_minimax_provider x2, test_deepseek_infra) - 'from src.models import Metadata' -> 'from src.type_aliases import Metadata' (test_project_manager_tracks, test_track_state_persistence, test_track_state_schema) Kept the 2 'import src.models as models' in test_provider_curation + test_providers_source_of_truth (they verify the backward-compat shim still re-exports PROVIDERS) and the 2 self-tests (test_models_no_top_level_pydantic + test_models_no_top_level_tomli_w) that exercise the shim's no-leak invariants.
88 lines
2.5 KiB
Python
88 lines
2.5 KiB
Python
import pytest
|
|
import asyncio
|
|
from src import ai_client
|
|
from src import mcp_client
|
|
from src.tool_presets import ToolPreset, Tool, Tool, ToolPreset
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_auto_approval():
|
|
# Setup a preset with read_file as auto
|
|
preset = ToolPreset(name="AutoTest", categories={
|
|
"General": [Tool(name="read_file", approval="auto")]
|
|
})
|
|
|
|
with patch("src.tool_presets.ToolPresetManager.load_all", return_value={"AutoTest": preset}):
|
|
ai_client.set_tool_preset("AutoTest")
|
|
|
|
# Mock mcp_client.async_dispatch to avoid actual file reads
|
|
with patch("src.mcp_client.async_dispatch", return_value="File Content") as mock_dispatch:
|
|
# pre_tool_callback should NOT be called
|
|
mock_cb = MagicMock()
|
|
|
|
name, call_id, out, _ = await ai_client._execute_single_tool_call_async(
|
|
name="read_file",
|
|
args={"path": "test.txt"},
|
|
call_id="call_1",
|
|
base_dir=".",
|
|
pre_tool_callback=mock_cb,
|
|
qa_callback=None,
|
|
r_idx=0
|
|
)
|
|
|
|
assert out == "File Content"
|
|
mock_cb.assert_not_called()
|
|
mock_dispatch.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_ask_approval():
|
|
# Setup a preset with run_powershell as ask
|
|
preset = ToolPreset(name="AskTest", categories={
|
|
"General": [Tool(name="run_powershell", approval="ask")]
|
|
})
|
|
|
|
with patch("src.tool_presets.ToolPresetManager.load_all", return_value={"AskTest": preset}):
|
|
ai_client.set_tool_preset("AskTest")
|
|
|
|
# pre_tool_callback SHOULD be called
|
|
mock_cb = MagicMock(return_value="Success")
|
|
|
|
name, call_id, out, _ = await ai_client._execute_single_tool_call_async(
|
|
name="run_powershell",
|
|
args={"script": "dir"},
|
|
call_id="call_2",
|
|
base_dir=".",
|
|
pre_tool_callback=mock_cb,
|
|
qa_callback=None,
|
|
r_idx=0
|
|
)
|
|
|
|
assert out == "Success"
|
|
mock_cb.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_rejection():
|
|
# Setup a preset with run_powershell as ask
|
|
preset = ToolPreset(name="AskTest", categories={
|
|
"General": [Tool(name="run_powershell", approval="ask")]
|
|
})
|
|
|
|
with patch("src.tool_presets.ToolPresetManager.load_all", return_value={"AskTest": preset}):
|
|
ai_client.set_tool_preset("AskTest")
|
|
|
|
# mock_cb returns None (rejected)
|
|
mock_cb = MagicMock(return_value=None)
|
|
|
|
name, call_id, out, _ = await ai_client._execute_single_tool_call_async(
|
|
name="run_powershell",
|
|
args={"script": "dir"},
|
|
call_id="call_3",
|
|
base_dir=".",
|
|
pre_tool_callback=mock_cb,
|
|
qa_callback=None,
|
|
r_idx=0
|
|
)
|
|
|
|
assert "USER REJECTED" in out
|
|
mock_cb.assert_called_once()
|