69 lines
1.9 KiB
Python
69 lines
1.9 KiB
Python
import json
|
|
from unittest.mock import patch, MagicMock
|
|
from agent.api_client import APIClient
|
|
|
|
|
|
def test_parse_json_valid():
|
|
text = '{"groups": {"g1": {"records": []}}}'
|
|
result = APIClient._parse_json(text)
|
|
assert result is not None
|
|
assert 'groups' in result
|
|
|
|
|
|
def test_parse_json_with_markdown_wrapper():
|
|
text = '```json\n{"key": "value"}\n```'
|
|
result = APIClient._parse_json(text)
|
|
assert result is not None
|
|
assert result['key'] == 'value'
|
|
|
|
|
|
def test_parse_json_with_text_before():
|
|
text = '说明\n{"key": "value"}\n以上'
|
|
result = APIClient._parse_json(text)
|
|
assert result is not None
|
|
assert result['key'] == 'value'
|
|
|
|
|
|
def test_parse_json_invalid():
|
|
text = '这不是有效的JSON。'
|
|
result = APIClient._parse_json(text)
|
|
assert result is None
|
|
|
|
|
|
@patch('agent.api_client.requests.post')
|
|
def test_generate_success(mock_post):
|
|
mock_response = MagicMock()
|
|
mock_response.json.return_value = {
|
|
'choices': [{'message': {'content': '{"groups": {"g1": {"type": "json"}}}'}}]
|
|
}
|
|
mock_response.raise_for_status = MagicMock()
|
|
mock_post.return_value = mock_response
|
|
|
|
client = APIClient(api_key='test-key')
|
|
result = client.generate("test prompt")
|
|
|
|
assert result['groups']['g1']['type'] == 'json'
|
|
|
|
|
|
@patch('agent.api_client.requests.post')
|
|
def test_generate_retry_on_json_error(mock_post):
|
|
bad_response = MagicMock()
|
|
bad_response.json.return_value = {
|
|
'choices': [{'message': {'content': 'invalid response'}}]
|
|
}
|
|
bad_response.raise_for_status = MagicMock()
|
|
|
|
good_response = MagicMock()
|
|
good_response.json.return_value = {
|
|
'choices': [{'message': {'content': '{"result": "ok"}'}}]
|
|
}
|
|
good_response.raise_for_status = MagicMock()
|
|
|
|
mock_post.side_effect = [bad_response, good_response]
|
|
|
|
client = APIClient(api_key='test-key', max_retries=3)
|
|
result = client.generate("test")
|
|
|
|
assert result['result'] == 'ok'
|
|
assert mock_post.call_count == 2
|