Codex
138 lines
6.1 KiB
Python
138 lines
6.1 KiB
Python
import unittest
|
|
|
|
from fastapi import HTTPException
|
|
|
|
from app.clients.shared_ai_platform import HybridSharedAIPlatformClient
|
|
|
|
|
|
class StubRemoteClient:
|
|
def __init__(self):
|
|
self.calls = []
|
|
self.raise_for = {}
|
|
|
|
def _run(self, method_name, *args):
|
|
self.calls.append((method_name, args))
|
|
maybe_error = self.raise_for.get(method_name)
|
|
if maybe_error is not None:
|
|
raise maybe_error
|
|
return {"source": "remote", "method": method_name}
|
|
|
|
def stream_answer(self, current_user, conversation_id, message, review_context=None):
|
|
self.calls.append(("stream_answer", (current_user, conversation_id, message, review_context)))
|
|
maybe_error = self.raise_for.get("stream_answer")
|
|
if maybe_error is not None:
|
|
raise maybe_error
|
|
yield 'data: {"type":"text","content":"remote"}\n\n'
|
|
|
|
def list_knowledge_bases(self, current_user, page, page_size):
|
|
return self._run("list_knowledge_bases", current_user, page, page_size)
|
|
|
|
def list_conversations(self, current_user, page, page_size):
|
|
return self._run("list_conversations", current_user, page, page_size)
|
|
|
|
def get_conversation_detail(self, current_user, conversation_id):
|
|
return self._run("get_conversation_detail", current_user, conversation_id)
|
|
|
|
def run_legal_review(self, current_user, payload):
|
|
return self._run("run_legal_review", current_user, payload)
|
|
|
|
def create_knowledge_base(self, current_user, name, description):
|
|
return self._run("create_knowledge_base", current_user, name, description)
|
|
|
|
def list_documents(self, current_user, knowledge_base_id, page, page_size):
|
|
return self._run("list_documents", current_user, knowledge_base_id, page, page_size)
|
|
|
|
def upload_document(self, current_user, knowledge_base_id, file_name, mime_type, content):
|
|
return self._run("upload_document", current_user, knowledge_base_id, file_name, mime_type, content)
|
|
|
|
def retry_document(self, current_user, knowledge_base_id, document_id, file_name=None, mime_type=None, content=None):
|
|
return self._run("retry_document", current_user, knowledge_base_id, document_id, file_name, mime_type, content)
|
|
|
|
|
|
class StubLocalClient:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def _run(self, method_name, *args):
|
|
self.calls.append((method_name, args))
|
|
return {"source": "local", "method": method_name}
|
|
|
|
def stream_answer(self, current_user, conversation_id, message, review_context=None):
|
|
self.calls.append(("stream_answer", (current_user, conversation_id, message, review_context)))
|
|
yield 'data: {"type":"text","content":"local"}\n\n'
|
|
|
|
def list_knowledge_bases(self, current_user, page, page_size):
|
|
return self._run("list_knowledge_bases", current_user, page, page_size)
|
|
|
|
def list_conversations(self, current_user, page, page_size):
|
|
return self._run("list_conversations", current_user, page, page_size)
|
|
|
|
def get_conversation_detail(self, current_user, conversation_id):
|
|
return self._run("get_conversation_detail", current_user, conversation_id)
|
|
|
|
def run_legal_review(self, current_user, payload):
|
|
return self._run("run_legal_review", current_user, payload)
|
|
|
|
def create_knowledge_base(self, current_user, name, description):
|
|
return self._run("create_knowledge_base", current_user, name, description)
|
|
|
|
def list_documents(self, current_user, knowledge_base_id, page, page_size):
|
|
return self._run("list_documents", current_user, knowledge_base_id, page, page_size)
|
|
|
|
def upload_document(self, current_user, knowledge_base_id, file_name, mime_type, content):
|
|
return self._run("upload_document", current_user, knowledge_base_id, file_name, mime_type, content)
|
|
|
|
def retry_document(self, current_user, knowledge_base_id, document_id, file_name=None, mime_type=None, content=None):
|
|
return self._run("retry_document", current_user, knowledge_base_id, document_id, file_name, mime_type, content)
|
|
|
|
|
|
class HybridSharedAIPlatformClientTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.remote = StubRemoteClient()
|
|
self.local = StubLocalClient()
|
|
self.client = HybridSharedAIPlatformClient(self.remote, self.local)
|
|
self.current_user = {"id": "user-1", "tenant_id": "tenant-1", "role": "admin"}
|
|
|
|
def test_read_calls_use_remote_when_remote_is_healthy(self):
|
|
payload = self.client.list_conversations(self.current_user, 1, 20)
|
|
|
|
self.assertEqual(payload["source"], "remote")
|
|
self.assertEqual(self.local.calls, [])
|
|
|
|
def test_read_calls_fallback_to_local_on_remote_500(self):
|
|
self.remote.raise_for["list_knowledge_bases"] = HTTPException(status_code=500, detail="remote failure")
|
|
|
|
payload = self.client.list_knowledge_bases(self.current_user, 1, 20)
|
|
|
|
self.assertEqual(payload["source"], "local")
|
|
self.assertEqual(self.local.calls[0][0], "list_knowledge_bases")
|
|
|
|
def test_read_calls_do_not_mask_remote_403(self):
|
|
self.remote.raise_for["get_conversation_detail"] = HTTPException(status_code=403, detail="forbidden")
|
|
|
|
with self.assertRaises(HTTPException) as context:
|
|
self.client.get_conversation_detail(self.current_user, "conv-1")
|
|
|
|
self.assertEqual(context.exception.status_code, 403)
|
|
self.assertEqual(self.local.calls, [])
|
|
|
|
def test_stream_answer_falls_back_to_local_when_remote_stream_crashes(self):
|
|
self.remote.raise_for["stream_answer"] = RuntimeError("stream offline")
|
|
|
|
chunks = list(self.client.stream_answer(self.current_user, None, "hello", None))
|
|
|
|
self.assertEqual(chunks, ['data: {"type":"text","content":"local"}\n\n'])
|
|
self.assertEqual(self.local.calls[0][0], "stream_answer")
|
|
|
|
def test_write_calls_fallback_to_local_on_remote_500(self):
|
|
self.remote.raise_for["create_knowledge_base"] = HTTPException(status_code=502, detail="bad gateway")
|
|
|
|
payload = self.client.create_knowledge_base(self.current_user, "测试知识库", "desc")
|
|
|
|
self.assertEqual(payload["source"], "local")
|
|
self.assertEqual(self.local.calls[0][0], "create_knowledge_base")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|