v1.0

tests/test_api_client.py

@@ -0,0 +1,68 @@
+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