115 lines
4.1 KiB
Python
115 lines
4.1 KiB
Python
from unittest.mock import patch, MagicMock
|
|
from src import ai_client
|
|
import json
|
|
import pytest
|
|
|
|
def test_deepseek_model_selection() -> None:
|
|
"""
|
|
Verifies that ai_client.set_provider('deepseek', 'deepseek-chat') correctly updates the internal state.
|
|
"""
|
|
ai_client.set_provider("deepseek", "deepseek-chat")
|
|
assert ai_client._provider == "deepseek"
|
|
assert ai_client._model == "deepseek-chat"
|
|
|
|
@patch("requests.post")
|
|
def test_deepseek_completion_logic(mock_post: MagicMock) -> None:
|
|
"""
|
|
Verifies that ai_client.send() correctly calls the DeepSeek API and returns content.
|
|
"""
|
|
ai_client.set_provider("deepseek", "deepseek-chat")
|
|
with patch("src.ai_client._load_credentials", return_value={"deepseek": {"api_key": "test-key"}}):
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"choices": [{"message": {"content": "Hello World"}, "finish_reason": "stop"}]
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
result = ai_client.send(md_content="Context", user_message="Hi", base_dir=".")
|
|
assert result == "Hello World"
|
|
assert mock_post.called
|
|
|
|
@patch("requests.post")
|
|
def test_deepseek_reasoning_logic(mock_post: MagicMock) -> None:
|
|
"""
|
|
Verifies that reasoning_content is captured and wrapped in <thinking> tags.
|
|
"""
|
|
ai_client.set_provider("deepseek", "deepseek-reasoner")
|
|
with patch("src.ai_client._load_credentials", return_value={"deepseek": {"api_key": "test-key"}}):
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"choices": [{
|
|
"message": {"content": "Final answer", "reasoning_content": "Chain of thought"},
|
|
"finish_reason": "stop"
|
|
}]
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
result = ai_client.send(md_content="Context", user_message="Hi", base_dir=".")
|
|
assert "<thinking>\nChain of thought\n</thinking>" in result
|
|
assert "Final answer" in result
|
|
|
|
@patch("requests.post")
|
|
def test_deepseek_tool_calling(mock_post: MagicMock) -> None:
|
|
"""
|
|
Verifies that DeepSeek provider correctly identifies and executes tool calls.
|
|
"""
|
|
ai_client.set_provider("deepseek", "deepseek-chat")
|
|
with patch("src.ai_client._load_credentials", return_value={"deepseek": {"api_key": "test-key"}}), \
|
|
patch("src.mcp_client.dispatch") as mock_dispatch:
|
|
|
|
# Round 1: Model calls a tool
|
|
mock_resp1 = MagicMock()
|
|
mock_resp1.status_code = 200
|
|
mock_resp1.json.return_value = {
|
|
"choices": [{
|
|
"message": {
|
|
"content": "I will read the file",
|
|
"tool_calls": [{
|
|
"id": "call_1",
|
|
"type": "function",
|
|
"function": {"name": "read_file", "arguments": '{"path": "test.txt"}'}
|
|
}]
|
|
},
|
|
"finish_reason": "tool_calls"
|
|
}]
|
|
}
|
|
|
|
# Round 2: Model provides final answer
|
|
mock_resp2 = MagicMock()
|
|
mock_resp2.status_code = 200
|
|
mock_resp2.json.return_value = {
|
|
"choices": [{"message": {"content": "File content is: Hello World"}, "finish_reason": "stop"}]
|
|
}
|
|
|
|
mock_post.side_effect = [mock_resp1, mock_resp2]
|
|
mock_dispatch.return_value = "Hello World"
|
|
|
|
result = ai_client.send(md_content="Context", user_message="Read test.txt", base_dir=".")
|
|
assert "File content is: Hello World" in result
|
|
assert mock_dispatch.called
|
|
mock_dispatch.assert_called_with("read_file", {"path": "test.txt"})
|
|
|
|
@patch("requests.post")
|
|
def test_deepseek_streaming(mock_post: MagicMock) -> None:
|
|
"""
|
|
Verifies that DeepSeek provider correctly aggregates streaming chunks.
|
|
"""
|
|
ai_client.set_provider("deepseek", "deepseek-chat")
|
|
with patch("src.ai_client._load_credentials", return_value={"deepseek": {"api_key": "test-key"}}):
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
|
|
# Mocking an iterable response for stream=True
|
|
chunks = [
|
|
'data: {"choices": [{"delta": {"content": "Hello "}}]}\n',
|
|
'data: {"choices": [{"delta": {"content": "World"}}]}\n',
|
|
'data: [DONE]\n'
|
|
]
|
|
mock_response.iter_lines.return_value = [c.encode('utf-8') for c in chunks]
|
|
mock_post.return_value = mock_response
|
|
|
|
result = ai_client.send(md_content="Context", user_message="Stream test", base_dir=".", stream=True)
|
|
assert result == "Hello World"
|