BossAi/test_api.py

from fastapi.testclient import TestClient
import unittest
from api import app
import api as api_module

client = TestClient(app)


class ApiTests(unittest.TestCase):
    def setUp(self):
        """每个测试前清理限流桶,避免跨测试影响。"""
        api_module._rate_buckets.clear()

    def test_health_and_ephemeral_flow(self):
        self.assertEqual(client.get('/health').status_code, 200)
        sid = client.post('/sessions', json={'mode': 'ephemeral'}).json()['session_id']
        r = client.post(f'/sessions/{sid}/ask', json={'question': '员工工伤后能否开除'})
        self.assertEqual(r.json()['risk']['level'], 'orange')
        self.assertTrue(client.post(f'/sessions/{sid}/destroy').json()['destroyed'])
        self.assertEqual(client.post(f'/sessions/{sid}/ask', json={'question': 'test'}).status_code, 410)

    def test_sensitive_file_limit(self):
        sid = client.post('/sessions', json={'mode': 'confidential'}).json()['session_id']
        r = client.post(f'/sessions/{sid}/file', files={'file': ('a.txt', b'hello')})
        self.assertEqual(r.status_code, 200)

    def test_text_file_and_knowledge_search(self):
        sid = client.post('/sessions', json={'mode': 'normal'}).json()['session_id']
        r = client.post(f'/sessions/{sid}/file', files={'file': ('a.txt', '付款违约条款'.encode())})
        self.assertIn('付款违约', r.json()['text_preview'])

    def test_knowledge_rag_search(self):
        """知识库 RAG 检索:中文 2-gram 分词"""
        results = client.get('/knowledge/search?q=工伤解除').json()['results']
        self.assertGreater(len(results), 0)
        self.assertIn('score', results[0])

    def test_knowledge_add_and_list(self):
        """知识库新增和列表"""
        r = client.post('/knowledge', json={
            'title': '测试法规', 'content': '测试内容', 'doc_type': 'law', 'region': '全国'
        })
        self.assertEqual(r.status_code, 200)
        self.assertIn('id', r.json())
        docs = client.get('/knowledge').json()['results']
        self.assertGreater(len(docs), 0)

    def test_settings_get_and_update(self):
        """设置持久化:读取和更新"""
        s = client.get('/settings').json()['settings']
        self.assertIn('cloud_call_allowed', s)
        r = client.put('/settings', json={'settings': {'long_term_memory': 'true'}})
        self.assertEqual(r.status_code, 200)
        self.assertEqual(r.json()['settings']['long_term_memory'], 'true')

    def test_expert_list_and_select(self):
        """专家列表和选定"""
        # 列出法务专家
        r = client.get('/experts?category=legal')
        self.assertEqual(r.status_code, 200)
        self.assertGreater(len(r.json()['results']), 0)
        expert_id = r.json()['results'][0]['id']
        # 专家详情
        detail = client.get(f'/experts/{expert_id}')
        self.assertEqual(detail.status_code, 200)
        self.assertIn('ratings', detail.json())

    def test_ask_with_rag_context(self):
        """法务场景 ask 接口注入 RAG 上下文"""
        sid = client.post('/sessions', json={'mode': 'confidential'}).json()['session_id']
        r = client.post(f'/sessions/{sid}/ask', json={
            'question': '员工工伤后能否开除', 'scene': 'legal'
        })
        self.assertEqual(r.status_code, 200)
        self.assertIn('rag_context', r.json())
        self.assertIn('inference_location', r.json())

    def test_file_parse_docx(self):
        """DOCX 文件解析"""
        import docx
        import io
        doc = docx.Document()
        doc.add_paragraph('测试合同条款')
        buf = io.BytesIO()
        doc.save(buf)
        sid = client.post('/sessions', json={'mode': 'confidential'}).json()['session_id']
        r = client.post(f'/sessions/{sid}/file', files={'file': ('test.docx', buf.getvalue())})
        self.assertEqual(r.status_code, 200)
        self.assertIn('测试合同', r.json()['text_preview'])
        self.assertGreater(r.json()['text_length'], 0)

    def test_vector_search(self):
        """向量检索:TF-IDF 余弦相似度"""
        r = client.get('/knowledge/vector-search?q=工伤解除')
        self.assertEqual(r.status_code, 200)
        self.assertGreater(len(r.json()['results']), 0)
        self.assertIn('score', r.json()['results'][0])

    def test_destroy_status(self):
        """销毁状态查询"""
        sid = client.post('/sessions', json={'mode': 'ephemeral'}).json()['session_id']
        client.post(f'/sessions/{sid}/destroy')
        r = client.get(f'/sessions/{sid}/destroy-status')
        self.assertEqual(r.status_code, 200)
        self.assertTrue(r.json()['destroyed'])
        self.assertGreater(len(r.json()['destroy_log']), 0)

    def test_action_summary_save(self):
        """脱敏行动摘要保存"""
        r = client.post('/summaries', json={
            'session_id': 'test',
            'summary': {'risk_level': 'orange', 'conclusion': '测试结论', 'actions': ['行动1']}
        })
        self.assertEqual(r.status_code, 200)
        self.assertIn('id', r.json())
        self.assertEqual(r.json()['summary']['risk_level'], 'orange')

    def test_metrics(self):
        """可观测性指标"""
        r = client.get('/metrics')
        self.assertEqual(r.status_code, 200)
        self.assertIn('request_count', r.json())
        self.assertIn('model_calls', r.json())

    def test_rate_limit(self):
        """限流:设置低阈值后快速请求应触发 429"""
        import importlib
        import api as api_module
        original = api_module.RATE_LIMIT
        api_module.RATE_LIMIT = 2  # 临时设为 2/秒
        try:
            codes = []
            for i in range(10):
                r = client.get('/settings')
                codes.append(r.status_code)
            self.assertIn(429, codes)
        finally:
            api_module.RATE_LIMIT = original
            # 清理限流桶,避免影响后续测试
            api_module._rate_buckets.clear()


if __name__ == '__main__':
    unittest.main()